project.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "dd2b25d5",
   "metadata": {},
   "source": [
    "# Thesis Project\n",
    "\n",
    "This is the software project for the bachelor thesis on \"Modelling of Perturbations in Seismocardiography Signals\". The complete source code is included in this notebook and contains the full processing pipeline, including the loading and pre-processing of the datasets, the training of the autoencoders, the evaluation for the four experiments and further utility functions to create visualisations of the results. The rest of the project is structured as follows: First, the necessary imports for the project are made, followed by the configuration of some important settings on which the rest of the project is based. Then the dataset(s) will be imported and pre-processed, followed by the definition, hyperparameter tuning and training of the autoencoders. Finally, the testing and evaluation for the four different experiments and their respective sub-experiments are specified, after which functions for the visualisation of the results are defined."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "157ff5a7",
   "metadata": {},
   "source": [
    "## Requirements and Imports\n",
    "\n",
    "All requirements in this project are managed by `Pipenv` and are specified in the `Pipfile` which is located in the same directory as this notebook. This code cell imports all the packages required for this project, including some basic imports such as numpy, pandas, matplotlib or other basic Python packages, all imports related to PyTorch and RayTune, and finally all imports related to either scipy or scikit-learn/scikit-image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b4b0800",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Basic imports\n",
    "import os\n",
    "import re\n",
    "import glob\n",
    "import wfdb\n",
    "import time\n",
    "import json\n",
    "import random\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "from wfdb import processing\n",
    "from functools import partial\n",
    "\n",
    "# PyTorch-related imports\n",
    "import torch\n",
    "import torchinfo\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "\n",
    "# RayTune-related imports\n",
    "from ray import tune, put, get, ray\n",
    "from ray.tune import CLIReporter\n",
    "from ray.tune.schedulers import ASHAScheduler\n",
    "\n",
    "# Scipy and Sklearn/Skimage-related imports\n",
    "from scipy import signal, stats\n",
    "from statsmodels.tsa.stattools import ccf\n",
    "from skimage.restoration import denoise_wavelet\n",
    "from sklearn.metrics import mean_squared_error, mean_absolute_error"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "0e8cb7b9",
   "metadata": {},
   "source": [
    "## Configuration\n",
    "\n",
    "The following code cell can be used to configure certain general aspects of the project related to random number generation, graphics card utilisation, machine learning model and data loading, noise modelling, and training. In particular, for the data loading settings, not only the directory from which the files are loaded is set, but also the sampling rate at which all the data is resampled, and settings for splitting the data with a given `window_size`, where the windows have an offset of a given `step_size`. To model the White Gaussian Noise (WGN), the noise distributions for testing and training are specified here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c145768a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set a random seed for pytorch (always split the data the same way), for the random package and for numpy. Required for reproducibility of the experiments.\n",
    "seed = 42\n",
    "random.seed(seed)\n",
    "np.random.seed(seed)\n",
    "torch.manual_seed(seed)\n",
    "torch.cuda.manual_seed(seed)\n",
    "\n",
    "# If a GPU (Cuda) is available, it will be used to train and test the machine learning models.\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "# The directory for saving and loading the trained autoencoder models.\n",
    "model_directory = './models/'\n",
    "\n",
    "# Here, some basic settings for the datasets and are set.\n",
    "data_directory = './data/' # Where to load the datasets from. Should be the same directory as in the `init.sh` script.\n",
    "sampling_rate = 100 # Hz # Resample the data according to this sampling rate.\n",
    "window_size = 512 # Of how many datapoints one sample should consist.\n",
    "step_size = 64 # How many datapoints are between the different samples. Implemented with a rolling window.\n",
    "\n",
    "train_prob = 0.8 # Relative size of the train set to complete dataset\n",
    "val_prob = 0.2 # Relative size of the validation set to the train set\n",
    "\n",
    "# The target noise SNR values for training and testing in dB.\n",
    "# train_target_snrs = [-2.5, 0, 2.5] # WGN1\n",
    "# train_target_snrs = [-2.5, 0, 2.5, 2.5, 5, 7.5] # WGN2\n",
    "train_target_snrs = [-2.5, 0, 2.5, 7.5, 10, 12.5] # WGN3\n",
    "# train_target_snrs = [-2.5, 0, 2.5, 17.5, 20, 22.5] # WGN4\n",
    "\n",
    "test_target_snrs =  [-1, 0, 2, 5, 7] # WGN5\n",
    "# test_target_snrs = [-1]\n",
    "# test_target_snrs = [7]\n",
    "\n",
    "# Number of epochs for training.\n",
    "num_epochs = 64"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "9fbeb9e6",
   "metadata": {},
   "source": [
    "## Dataset\n",
    "\n",
    "The first code cell contains some functions for normalising and denormalising a one- or two-dimensional array, but also a Butterworth bandpass filter with a standard configuration of order 4 and cut-off frequencies of 5 and 30 Hz for SCG signals. In addition, a filter is implemented to remove peaks from a given signal by calculating the z-score of each value, where any datapoint with a value above a given threshold is set to the mean of the signal.\n",
    "\n",
    "The second cell includes a utility class used to manage the data required for training and testing. It inherits from the PyTorch `Dataset` class, which provides tools such as splitting the data into training, validation and test subsets, or using it in dataloaders to efficiently access the data in minibatches during training and testing. The data for the experiments may consist the [CEBS](https://physionet.org/content/cebsdb/1.0.0/) dataset, the [IEEE](https://ieee-dataport.org/documents/mechanocardiograms-ecg-reference) dataset and/or another dataset called [ECMS](https://zenodo.org/record/5279448), as described in the methodology of the experiments in the thesis. Either clean SCG or clean ECG samples can be used as y-samples and SCG or noisy SCG samples as x-samples. \n",
    "\n",
    "The data from the specified datasets is loaded directly when a `SCGData` object is initialised, depending on the specific configuration of that object, which is described below. Each dataset is imported according to its specific characteristics, the schemes are normalised, the dataset is resampled per subject and a fourth order Butterworth bandpass filter with cut-off frequencies of 5 and 30 Hz is applied to the SCG samples, followed by a filter based on the calculation of the z-score of each datapoint to remove outliers with a threshold of 7. Lastly, the data per subject is normalised between -1 and 1 and then split into samples with a size of `window_size` and an offset of `step_size` using a rolling window. All datasets are then combined. \n",
    "\n",
    "Initially, no noise is added to the data so that training and testing can be done with different noise levels. The noise generated is WGN and/or a second noise model focussing on motion noise. The two noise models can be used independently or together, the latter with a certain probability of generating WGN. The motion-focused noise model takes motion data from a dataset that includes [data from a chest-worn accelerometer](https://archive.ics.uci.edu/ml/datasets/Activity+Recognition+from+Single+Chest-Mounted+Accelerometer) during 7 different activities. The gravity component is removed from the motion data as it is already present in the SCG datasets and outliers are removed using the z-score method with a threshold of 5. To model the motion noise, `window_size` long random consecutive samples are extracted and 80% of their values are added to the clean SCG samples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46e9a7aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "A method to normalise the input signal between -1 and 1. Returns the minimum and maximum values of the original signal as well for later denormalisation.\n",
    ":param data: The signal (array) that should be normalised.\n",
    "'''\n",
    "def normalise(data):\n",
    "    # If the data consists of several arrays, concatenate them to find the minimum and maximum.\n",
    "    data_concat = np.concatenate(data) if data[0].ndim == 1 else data\n",
    "    min_val = np.min(data_concat)\n",
    "    max_val = np.max(data_concat)\n",
    "    data_norm = 2 * (data - min_val) / (max_val - min_val) - 1\n",
    "    return data_norm, min_val, max_val\n",
    "\n",
    "'''\n",
    "Denormalises the input signal to its original form, given the minimum and maximum values of the original signal.\n",
    ":param data: The signal (array) that should be denormalised.\n",
    ":param min_val: The minimum value of the original, not normalised signal.\n",
    ":param max_val: The maximum value of the original, not normalised signal.\n",
    "'''\n",
    "def denormalise(data, min_val, max_val):\n",
    "    data_denorm = (max_val * data + max_val - data * min_val + min_val) / 2\n",
    "    return data_denorm\n",
    "\n",
    "'''\n",
    "A Butterworth bandpass filter for the input signal.\n",
    ":param input_signal: The signal to which the filter should be applied to.\n",
    ":param order: The order of the filter.\n",
    ":param fs: The frequency space.\n",
    ":param lowcut: The highcut of the filter.\n",
    ":param highcut: The lowcut of the filter.\n",
    "'''\n",
    "def bandpass(input_signal, order=4, fs=100.0, lowcut=5.0, highcut=30.0):\n",
    "    nyq = 0.5 * fs\n",
    "    low = lowcut / nyq\n",
    "    high = highcut / nyq\n",
    "\n",
    "    b, a, = signal.butter(order, [low, high], 'bandpass', analog=False)\n",
    "    y = signal.filtfilt(b, a, input_signal, axis=0)\n",
    "\n",
    "    return y\n",
    "\n",
    "'''\n",
    "A method to filter out peaks of a given input signal by comparing the z-score of each datapoint to a threshold value. \n",
    ":param input_signal: The signal to which the filter should be applied to.\n",
    ":param z_score_th: Each datapoint with a z-score above this threshold will be set to the mean of the signal.\n",
    "'''\n",
    "def filter_peaks(input_signal, z_score_th=7):\n",
    "    output_signal = input_signal.copy()\n",
    "    output_signal.loc[np.abs(stats.zscore(output_signal)) > z_score_th] = output_signal.mean()\n",
    "    return output_signal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9c8d52a",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SCGData(Dataset):\n",
    "\n",
    "    '''\n",
    "    A dataset class to manage all the different datasets being used in this project.\n",
    "\n",
    "    :param directory: The directory where the datasets (each in its own subdirectory named 'CEBS', 'ECMS' or 'IEEE') are located.\n",
    "    :param datasets: A list of strings from ['CEBS', 'ECMS', 'IEEE'] of the datasets that should be loaded.\n",
    "    :param sampling_rate: The sampling rate at which the data from all datasets will be resampled per subject.\n",
    "    :param window_size: The size of the windows into which the data is divided per subject.\n",
    "    :param step_size: The number of datapoints between the respective windows. Implemented as a rolling window.\n",
    "    :param transform: The transformation as its own class that will be applied to the data samples when they are accessed.\n",
    "    :param ecg: If true, the y_data will be ECG data and otherwise the clean SCG data.\n",
    "    :param noise: If true, noise may be added to the clean SCG samples (`self.scg_clean`) by calling the `add_noise` method. Ultimately, this will be the x_data. \n",
    "                  Otherwise, calling this method will have no effects.\n",
    "    :param remove_cebs: A list of strings containing the filenames of the CEBS dataset that should be omitted. \n",
    "                        None, if no files should be omitted.\n",
    "    :param leave_cebs: A list of strings containing the filenames of the CEBS dataset that should be included. All other files will be omitted. \n",
    "                       None if all files should be included.\n",
    "    '''\n",
    "    def __init__(self, directory, datasets, sampling_rate, window_size, step_size, transform=None, ecg=False, noise=True, remove_cebs=None, leave_cebs=None):\n",
    "\n",
    "        self.directory = directory\n",
    "        self.datasets = datasets\n",
    "\n",
    "        self.sampling_rate = sampling_rate\n",
    "        self.window_size = window_size\n",
    "        self.step_size = step_size\n",
    "        \n",
    "        self.transform = transform\n",
    "        \n",
    "        self.ecg = ecg\n",
    "        self.noise = noise\n",
    "        self.leave_cebs = leave_cebs\n",
    "        self.remove_cebs = remove_cebs\n",
    "\n",
    "        self.data_motion_noise = None # Saves the data loaded from the motion dataset to add motion noise later.\n",
    "        \n",
    "        # The min/max values after normalising the noisy x_data. Used for later denormalisation to calculate the SNR.\n",
    "        self.min_val = None\n",
    "        self.max_val = None\n",
    "\n",
    "        # Load the motion noise data.\n",
    "        self._load_motion_noise_data()\n",
    "\n",
    "        # Load and pre-process the requested datasets from `self.datasets`.\n",
    "        data = self._load_datasets()\n",
    "        \n",
    "        # Clean SCG data to add noise to it later in order to create the possibly noisy x_data.\n",
    "        self.scg_clean = data.SCG.to_numpy(copy=True)\n",
    "        \n",
    "        # The y values are either ECG samples (if ecg=True) or SCG samples (if ecg=False). Noise will be added later to the x_data.\n",
    "        self.x_data = data.SCG.to_numpy()\n",
    "        self.y_data = data.ECG.to_numpy() if self.ecg else data.SCG.to_numpy()\n",
    "\n",
    "        self.n_samples = data.shape[0]\n",
    "\n",
    "    '''\n",
    "    The function may be called to add noise to `self.x_data` when `self.noise` is True.\n",
    "\n",
    "    :param target_snrs: A list of SNRs to be added with the WGN model. \n",
    "    :param noise_model: This parameter can either be 'WGN' for the White Gaussian Noise, 'MOTION' for the motion-based noise model or 'BOTH' for both noise models to be added to the dataset.\n",
    "    :param wgn_prob: Specifies the probability that the WGN model will be added to the dataset when 'BOTH' is selected as the noise model.\n",
    "    :param use_bandpass: If true, a bandpass filter is applied to the noisy data after adding noise within the `add_noise` method.\n",
    "    '''\n",
    "    def add_noise(self, target_snrs, noise_model='WGN', wgn_prob=0.4, use_bandpass=False):\n",
    "\n",
    "        if self.noise:\n",
    "            if noise_model == 'WGN':\n",
    "                data_noise = np.array([self._add_gaussian(inputs, target_snrs) for inputs in self.scg_clean])\n",
    "            elif noise_model == 'MOTION':\n",
    "                data_noise = np.array([self._add_motion(inputs) for inputs in self.scg_clean])\n",
    "            elif noise_model == 'BOTH':\n",
    "                data_noise = np.array([self._add_gaussian(inputs, target_snrs) if random.random() < wgn_prob else self._add_motion(inputs) for inputs in self.scg_clean])\n",
    "\n",
    "            self.x_data, self.min_val, self.max_val = normalise(data_noise)\n",
    "\n",
    "            if use_bandpass:\n",
    "                self.x_data = np.array([bandpass(inputs, order=4, fs=100.0, lowcut=5.0, highcut=30.0) for inputs in self.x_data])\n",
    "\n",
    "    '''\n",
    "    A function to add motion noise to its inputs.\n",
    "\n",
    "    :param data: The data, to which the noise should be added.\n",
    "    '''\n",
    "    def _add_motion(self, data):\n",
    "        \n",
    "        # Randomly select `len(data)` consecutive datapoints from the processed noise dataset.\n",
    "        len_sample = len(data)\n",
    "        rows = range(self.data_motion_noise.shape[0])\n",
    "        index_start = random.randint(rows.start, rows.stop - len_sample)\n",
    "        noise = self.data_motion_noise.iloc[index_start:index_start + len_sample]\n",
    "        \n",
    "        # Only add 80 % of noise to the data.\n",
    "        return data + 0.8 * noise\n",
    "\n",
    "    '''\n",
    "    A function to add WGN to its input with specific SNRs.\n",
    "\n",
    "    :param data: The data to which the noise should be added.\n",
    "    :param target_snrs: A list of SNRs to be added. One SNR will be chosen randomly from this list. \n",
    "    '''\n",
    "    def _add_gaussian(self, data, target_snrs):\n",
    "\n",
    "        target_snr = random.choice(target_snrs)\n",
    "\n",
    "        data_watts_mean = np.mean(data ** 2)\n",
    "        data_db_mean = 10 * np.log10(data_watts_mean)\n",
    "\n",
    "        noise_db_mean = data_db_mean - target_snr\n",
    "        noise_watts_mean = 10 ** (noise_db_mean / 10)\n",
    "\n",
    "        mean_noise = 0\n",
    "        noise = np.random.normal(mean_noise, np.sqrt(noise_watts_mean), len(data))\n",
    "        \n",
    "        return data + noise\n",
    "\n",
    "    '''\n",
    "    Loads the motion noise dataset, upsamples, filters and normalises it.\n",
    "    '''\n",
    "    def _load_motion_noise_data(self):\n",
    "        \n",
    "        data = pd.DataFrame()\n",
    "\n",
    "        # Load the data from the 16 files\n",
    "        for i in range(1, 16):\n",
    "            rel_path = os.path.join('NOISE_ACC', f'{i}.csv')\n",
    "            df = pd.read_csv(os.path.join(self.directory, rel_path), header=None, usecols=[3, 4], names=['Z', 'labels'])\n",
    "            \n",
    "            df_resample = pd.DataFrame()\n",
    "            frequency = 52\n",
    "\n",
    "            # Upsample the z-axis data, subtract the mean (gravity component), remove outliers and normalise it \n",
    "            df_resample['Z'] = processing.resample_sig(df['Z'], frequency, self.sampling_rate)[0]\n",
    "            df_resample['Z'] = df_resample['Z'] - df_resample['Z'].mean() # Remove the gravity component from the acceleration data\n",
    "            df_resample['Z'] = filter_peaks(df_resample['Z'], z_score_th=5)\n",
    "            df_resample['Z'], _, _ = normalise(df_resample['Z'])\n",
    "\n",
    "            # Upsample the labels, round them and prevent values larger than 7 and smaller than 1.\n",
    "            df_resample['labels'] = processing.resample_sig(df['labels'], frequency, self.sampling_rate)[0]\n",
    "            df_resample['labels'] = df_resample['labels'].round(decimals=0)\n",
    "            df_resample.loc[df_resample['labels'] > 7, 'labels'] = 7\n",
    "            df_resample.loc[df_resample['labels'] < 1, 'labels'] = 1\n",
    "\n",
    "            data = pd.concat((data, df_resample))\n",
    "\n",
    "        # The different activities in the dataset.\n",
    "            # 1: Working at Computer\n",
    "            # 2: Standing Up, Walking and Going up\\down stairs\n",
    "            # 3: Standing\n",
    "            # 4: Walking\n",
    "            # 5: Going Up\\Down Stairs\n",
    "            # 6: Walking and Talking with Someone\n",
    "            # 7: Talking while Standing\n",
    "\n",
    "        # Here it is possible to only allow certain activities in the noise dataset.\n",
    "        # self.data_motion_noise = data.loc[data['labels'] == 4, 'Z']\n",
    "        self.data_motion_noise = data['Z']\n",
    "\n",
    "    \n",
    "    '''\n",
    "    Loads one file of the CEBS dataset and processes it. This includes resampling, filtering and splitting into groups of `self.window_size` long samples.\n",
    "\n",
    "    :param filename: The name of the specific file to be loaded.\n",
    "    '''\n",
    "    def _load_cebs(self, filename):\n",
    "        \n",
    "        # Using WFDB to import the files.\n",
    "        record = wfdb.rdsamp(filename)\n",
    "        data = record[0]\n",
    "        metadata = record[1]\n",
    "\n",
    "        frequency = metadata['fs']\n",
    "        column_names = metadata['sig_name']\n",
    "\n",
    "        data = pd.DataFrame(data, columns=column_names)\n",
    "        data.drop(['II', 'RESP'], axis=1, inplace=True)\n",
    "        data.rename(columns={'I': 'ECG'}, inplace=True)\n",
    "\n",
    "        # Resample the data into the desired sampling rate and use a bandpass filter on the SCG samples.\n",
    "        scg = processing.resample_sig(data['SCG'], frequency, self.sampling_rate)[0]\n",
    "        scg = bandpass(scg, order=4, fs=100.0, lowcut=5.0, highcut=30.0)\n",
    "        ecg = processing.resample_sig(data['ECG'], frequency, self.sampling_rate)[0]\n",
    "\n",
    "        data = pd.DataFrame({'SCG': scg, 'ECG': ecg})\n",
    "\n",
    "        # Filter out outliers from the SCG data\n",
    "        data['SCG'] = filter_peaks(data['SCG'], z_score_th=7)\n",
    "\n",
    "        # Split the data into groups of `self.window_size` datapoints and an offset of `self.step_size` datapoints.\n",
    "        data = pd.DataFrame([[x.ECG.to_numpy(), x.SCG.to_numpy()] for x in data.rolling(window=self.window_size, step=self.step_size) if x.shape[0] == self.window_size])\n",
    "        data.columns = ['ECG', 'SCG']\n",
    "\n",
    "        # Normalise the ECG and/or SCG data\n",
    "        # data['ECG'], _, _ = normalise(data['ECG'])\n",
    "        data['SCG'], _, _ = normalise(data['SCG'])\n",
    "        \n",
    "        return data\n",
    "\n",
    "\n",
    "    '''\n",
    "    Loads one file of the ECMS dataset and processes it. This includes resampling, filtering and splitting into groups of `self.window_size` long samples.\n",
    "\n",
    "    :param filename: The name of the specific file to be loaded.\n",
    "    '''\n",
    "    def _load_ecms(self, filename):\n",
    "\n",
    "        # Since the files are not structured the same, they have to be imported differently. The following code does that.\n",
    "        # Regex to identify files with specific names, as they are structured differently.\n",
    "        regex_up_1 = re.compile(r'./data/ECMS/UP-(((0|1)[0-9])|(20|21))-*') # Files named 'UP-[01-21]-Raw.csv'\n",
    "        regex_up_2 = re.compile(r'./data/ECMS/UP-(22|23)-*') # Files named 'UP-[22-23]-Raw.csv'\n",
    "\n",
    "        # Import the file depending on its structure.\n",
    "        if re.match(r'./data/ECMS/CP-*', filename):\n",
    "            data = pd.read_csv(filename, \n",
    "                    sep=',', \n",
    "                    header=1, \n",
    "                    skiprows=[2],\n",
    "                    usecols=['Shimmer_D0CD_ECG_LA-RA_24BIT_CAL', 'Shimmer_D0CD_Accel_LN_Z_CAL'], \n",
    "                    dtype={'Shimmer_D0CD_ECG_LA-RA_24BIT_CAL': 'float', 'Shimmer_D0CD_Accel_LN_Z_CAL': 'float'})\n",
    "            data.rename(columns={'Shimmer_D0CD_ECG_LA-RA_24BIT_CAL': 'ECG', 'Shimmer_D0CD_Accel_LN_Z_CAL': 'SCG'}, inplace=True)\n",
    "            frequency = 256\n",
    "        elif regex_up_1.match(filename):\n",
    "            data = pd.read_csv(filename, \n",
    "                    sep='\t', \n",
    "                    header=1, \n",
    "                    skiprows=[2],\n",
    "                    usecols=['ECG_ECG_LA-RA_24BIT_CAL', 'ECG_Accel_LN_Z_CAL'], \n",
    "                    dtype={'ECG_ECG_LA-RA_24BIT_CAL': 'float', 'ECG_Accel_LN_Z_CAL': 'float'})\n",
    "            data.rename(columns={'ECG_ECG_LA-RA_24BIT_CAL': 'ECG', 'ECG_Accel_LN_Z_CAL': 'SCG'}, inplace=True)\n",
    "            frequency = 256\n",
    "        elif regex_up_2.match(filename):\n",
    "            data = pd.read_csv(filename, \n",
    "                    sep=',', \n",
    "                    header=1, \n",
    "                    skiprows=[2],\n",
    "                    usecols=['ECG_ECG_LA-RA_24BIT_CAL', 'ECG_Accel_LN_Z_CAL'], \n",
    "                    dtype={'ECG_ECG_LA-RA_24BIT_CAL': 'float', 'ECG_Accel_LN_Z_CAL': 'float'})\n",
    "            data.rename(columns={'ECG_ECG_LA-RA_24BIT_CAL': 'ECG', 'ECG_Accel_LN_Z_CAL': 'SCG'}, inplace=True)\n",
    "            frequency = 512\n",
    "        elif re.match(r'./data/ECMS/UP-*', filename):\n",
    "            data = pd.read_csv(filename, \n",
    "                    sep=',', \n",
    "                    header=0, \n",
    "                    skiprows=[1],\n",
    "                    usecols=['ECG_ECG_LA-RA_24BIT_CAL', 'ECG_Accel_LN_Z_CAL'], \n",
    "                    dtype={'ECG_ECG_LA-RA_24BIT_CAL': 'float', 'ECG_Accel_LN_Z_CAL': 'float'})\n",
    "            data.rename(columns={'ECG_ECG_LA-RA_24BIT_CAL': 'ECG', 'ECG_Accel_LN_Z_CAL': 'SCG'}, inplace=True)            \n",
    "            frequency = 512\n",
    "        else:\n",
    "            return pd.DataFrame()\n",
    "\n",
    "        # Resample the data into the desired sampling rate and apply a bandpass filter to the SCG samples.\n",
    "        scg = processing.resample_sig(data['SCG'], frequency, self.sampling_rate)[0]\n",
    "        scg = bandpass(scg, order=4, fs=100.0, lowcut=5.0, highcut=30.0)\n",
    "        ecg = processing.resample_sig(data['ECG'], frequency, self.sampling_rate)[0]\n",
    "        # ecg = bandpass(ecg, order=4, fs=100.0, lowcut=8.0, highcut=20.0)\n",
    "\n",
    "        data = pd.DataFrame({'SCG': scg, 'ECG': ecg})\n",
    "        \n",
    "        # Filter out the outliers from the SCG data\n",
    "        data['SCG'] = filter_peaks(data['SCG'], z_score_th=7)\n",
    "        # data['ECG'] = filter_peaks(data['ECG'], z_score_th=7)\n",
    "\n",
    "        # Split the data into groups of `self.window_size` datapoints and an offset of `self.step_size` datapoints.\n",
    "        data = pd.DataFrame([[x.ECG.to_numpy(), x.SCG.to_numpy()] for x in data.rolling(window=self.window_size, step=self.step_size) if x.shape[0] == self.window_size])\n",
    "        data.columns = ['ECG', 'SCG']\n",
    "\n",
    "        # Normalise the ECG and/or SCG data\n",
    "        data['SCG'], _, _ = normalise(data['SCG'])\n",
    "        # data['ECG'], _, _ = normalise(data['ECG'])\n",
    "\n",
    "        return data\n",
    "\n",
    "    \n",
    "    '''\n",
    "    Loads one file of the IEEE dataset and processes it. This includes resampling, filtering and splitting into groups of `self.window_size` long samples.\n",
    "\n",
    "    :param filename: The name of the specific file to be loaded.\n",
    "    '''\n",
    "    def _load_ieee(self, filename):\n",
    "        \n",
    "        # Because of different headers in the files, a varying amount of rows has to be skipped.\n",
    "        if any(str(num) in filename for num in range(10,30)):\n",
    "            skiprows = 19\n",
    "        elif any(str(num) in filename for num in [8, 9]):\n",
    "            skiprows = 17\n",
    "        elif any(str(num) in filename for num in [2, 4, 5, 6, 7, 8]):\n",
    "            skiprows = 16\n",
    "        else:\n",
    "            skiprows = 15\n",
    "\n",
    "        data = pd.read_csv(filename, sep=' ', header=None, skiprows=skiprows, usecols=[0,3], names=['ECG', 'SCG'], dtype={'ECG': 'float', 'SCG': 'float'})\n",
    "        frequency = 800\n",
    "\n",
    "        # Resample the data into the desired sampling rate and apply a bandpass filter to the SCG samples.\n",
    "        scg = processing.resample_sig(data['SCG'], frequency, self.sampling_rate)[0]\n",
    "        scg = bandpass(scg, order=4, fs=100.0, lowcut=5.0, highcut=30.0)\n",
    "        ecg = processing.resample_sig(data['ECG'], frequency, self.sampling_rate)[0]\n",
    "        # ecg = bandpass(ecg, order=4, fs=100.0, lowcut=8.0, highcut=20.0)\n",
    "\n",
    "        data = pd.DataFrame({'SCG': scg, 'ECG': ecg})\n",
    "        \n",
    "        # Filter out the outliers from the SCG data\n",
    "        data['SCG'] = filter_peaks(data['SCG'], z_score_th=7)\n",
    "        # data['ECG'] = filter_peaks(data['ECG'], z_score_th=7)\n",
    "\n",
    "        # Split the data into groups of `self.window_size` datapoints and an offset of `self.step_size` datapoints.\n",
    "        data = pd.DataFrame([[x.ECG.to_numpy(), x.SCG.to_numpy()] for x in data.rolling(window=self.window_size, step=self.step_size) if x.shape[0] == self.window_size])\n",
    "        data.columns = ['ECG', 'SCG']\n",
    "        \n",
    "        # Normalise the ECG and.or SCG data\n",
    "        data['SCG'], _, _ = normalise(data['SCG'])\n",
    "        # data['ECG'], _, _ = normalise(data['ECG'])\n",
    "\n",
    "        return data\n",
    "\n",
    "\n",
    "    '''\n",
    "    Loads the datasets specified in `self.datasets`.\n",
    "    '''\n",
    "    def _load_datasets(self):\n",
    "\n",
    "        data = pd.DataFrame(dtype='object')\n",
    "\n",
    "        # Loads the CEBS dataset from the 'CEBS' directory in the directory folder.\n",
    "        if 'CEBS' in self.datasets:\n",
    "            print('Loading CEBS')\n",
    "            \n",
    "            directory = os.path.join(self.directory, 'CEBS')\n",
    "            if self.leave_cebs:\n",
    "                # Only load the specified files in `self.leave_cebs` if it is set\n",
    "                filenames = [os.path.join(directory, filename) for filename in self.leave_cebs]\n",
    "            else:\n",
    "                # Load all files otherwise\n",
    "                filenames = list(dict.fromkeys([x[:-4] for x in glob.glob(f'{directory}/*[0-9][0-9][0-9].*')]))\n",
    "            if self.remove_cebs:\n",
    "                # Remove the files specified in `self.remove_cebs`, if any.\n",
    "                for filename in self.remove_cebs:\n",
    "                    filenames.remove(os.path.join(directory, filename))\n",
    "            \n",
    "            data_cebs = pd.concat((self._load_cebs(filename) for filename in filenames), ignore_index=True)\n",
    "            data = pd.concat((data, data_cebs))\n",
    "            print(f'\\tLoaded {data_cebs.shape[0]} samples')\n",
    "\n",
    "        # Loads the ECMS dataset from the 'ECMS' directory in the directory folder.\n",
    "        if 'ECMS' in self.datasets:\n",
    "            print('Loading ECMS')\n",
    "            \n",
    "            directory = os.path.join(self.directory, 'ECMS')\n",
    "            filenames = glob.glob(f'{directory}/*')\n",
    "            data_ecms = pd.concat((self._load_ecms(filename) for filename in filenames), ignore_index=True)\n",
    "            data = pd.concat((data, data_ecms))\n",
    "            print(f'\\tLoaded {data_ecms.shape[0]} samples')\n",
    "\n",
    "        # Loads the IEEE dataset from the 'IEEE' directory in the directory folder.\n",
    "        if 'IEEE' in self.datasets:\n",
    "            print('Loading IEEE')\n",
    "            \n",
    "            directory = os.path.join(self.directory, 'IEEE')\n",
    "            filenames = glob.glob(f'{directory}/*')\n",
    "            data_ieee = pd.concat((self._load_ieee(filename) for filename in filenames), ignore_index=True)\n",
    "            data = pd.concat((data, data_ieee))\n",
    "            print(f'\\tLoaded {data_ieee.shape[0]} samples')\n",
    "        \n",
    "        print(f'\\nLoaded {data.shape[0]} samples in total')\n",
    "\n",
    "        return data\n",
    "\n",
    "    '''\n",
    "    Returns a sample with the specified index and applies the transformation, if it is set.\n",
    "\n",
    "    :param index: The index of the sample to be returned.\n",
    "    '''\n",
    "    def __getitem__(self, index):\n",
    "        sample = self.x_data[index].copy(), self.y_data[index].copy()\n",
    "\n",
    "        if self.transform:\n",
    "            sample = self.transform(sample)\n",
    "        \n",
    "        return sample\n",
    "    \n",
    "    '''\n",
    "    Returns the length of the dataset.\n",
    "    '''\n",
    "    def __len__(self):\n",
    "        return self.n_samples\n",
    "\n",
    "\n",
    "class ToTensor:\n",
    "\n",
    "    '''\n",
    "    Transformation class used to transform numpy arrays to PyTorch tensors.\n",
    "\n",
    "    :param sample: The sample from the dataset to be transformed.\n",
    "    '''\n",
    "    def __call__(self, sample):\n",
    "        inputs, targets = sample\n",
    "        return torch.from_numpy(inputs), torch.from_numpy(targets)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1492b678",
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "Loads the data and returns randomly splitted train and test datasets with a predetermined train size. \n",
    "The parameters are equal to those used in the __init__ method of the SCGData class.\n",
    "'''\n",
    "def load_data(*args, **kwargs):\n",
    "    dataset = SCGData(*args, **kwargs)\n",
    "    train_size = int(train_prob * len(dataset))\n",
    "    test_size = len(dataset) - train_size\n",
    "    return torch.utils.data.random_split(dataset, [train_size, test_size])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97834ec1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Imports the dataset and splits it into a train and test subset.\n",
    "# Omitted subjects for the last part of the third experiment: ['b001', 'm001', 'p001', 'b002', 'm002', 'p002']\n",
    "train_dataset, test_dataset = load_data(data_directory, ['CEBS', 'ECMS', 'IEEE'], sampling_rate, window_size, step_size, transform=ToTensor(), ecg=False, noise=True, remove_cebs=None, leave_cebs=None)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "b837dfa1",
   "metadata": {},
   "source": [
    "## Machine Learning Models\n",
    "\n",
    "The next two code cells implement 2 fully convolutional denoising autoencoders. \n",
    "\n",
    "The first one consists of 4 convolutional layers in the encoder and 4 transpose convolutional layers plus another convolutional output layer in the decoder. Each (transpose) convolutional layer except the output layer is followed by a batch normalisation layer, and ELU is used as the activation function, with the hyperbolic tangent function used for the output layer. The kernel size is set to 7 with a stride of 1 while the number of kernels, `kernel_count`, is halved with each convolutional layer in the encoder and is doubled with each layer in the decoder. The last layer in the encoder and decoder but also the output layer have only one kernel. No padding is applied to any layer except the output layer, where padding is set to `same`.\n",
    "\n",
    "The second autoencoder is nearly equivalent to the first one. However, the stride of the second last layer of the encoder and the second layer of the decoder are set to 2, while the output padding for the latter layer is set to 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba01cc41",
   "metadata": {},
   "outputs": [],
   "source": [
    "class AE1(nn.Module):\n",
    "    def __init__(self, kernel_count=128):\n",
    "        super(AE1, self).__init__()\n",
    "\n",
    "        self.encoder = nn.Sequential(\n",
    "            nn.Conv1d(1, int(kernel_count), 7, stride=1, padding=0),\n",
    "            nn.BatchNorm1d(int(kernel_count)),\n",
    "            nn.ELU(),\n",
    "            nn.Conv1d(int(kernel_count), int(kernel_count/2), 7, stride=1, padding=0),\n",
    "            nn.BatchNorm1d(int(kernel_count/2)),\n",
    "            nn.ELU(),\n",
    "            nn.Conv1d(int(kernel_count/2), int(kernel_count/4), 7, stride=1, padding=0),\n",
    "            nn.BatchNorm1d(int(kernel_count/4)),\n",
    "            nn.ELU(),\n",
    "            nn.Conv1d(int(kernel_count/4), 1, 7, stride=1, padding=0),\n",
    "            nn.BatchNorm1d(1),\n",
    "            nn.ELU()\n",
    "        )\n",
    "\n",
    "        self.decoder = nn.Sequential(\n",
    "            nn.ConvTranspose1d(1, int(kernel_count/4), 7, stride=1, padding=0),\n",
    "            nn.BatchNorm1d(int(kernel_count/4)),\n",
    "            nn.ELU(),\n",
    "            nn.ConvTranspose1d(int(kernel_count/4), int(kernel_count/2), 7, stride=1, padding=0),\n",
    "            nn.BatchNorm1d(int(kernel_count/2)),\n",
    "            nn.ELU(),\n",
    "            nn.ConvTranspose1d(int(kernel_count/2), int(kernel_count), 7, stride=1, padding=0),\n",
    "            nn.BatchNorm1d(int(kernel_count)),\n",
    "            nn.ELU(),\n",
    "            nn.ConvTranspose1d(int(kernel_count), 1, 7, stride=1, padding=0), \n",
    "            nn.BatchNorm1d(1),\n",
    "            nn.ELU(),\n",
    "            nn.Conv1d(1, 1, 7, stride=1, padding='same'),\n",
    "            nn.Tanh()\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        encoded = self.encoder(x)\n",
    "        decoded = self.decoder(encoded)\n",
    "        return decoded\n",
    "\n",
    "# Prints the structure of the autoencoder\n",
    "torchinfo.summary(AE1(), (8, 1, window_size), col_names=('input_size', 'output_size'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "031c1a44",
   "metadata": {},
   "outputs": [],
   "source": [
    "class AE2(nn.Module):\n",
    "    def __init__(self, kernel_count=128):\n",
    "        super(AE2, self).__init__()\n",
    "\n",
    "        self.encoder = nn.Sequential(\n",
    "            nn.Conv1d(1, int(kernel_count), 7, stride=1, padding=0),\n",
    "            nn.BatchNorm1d(int(kernel_count)),\n",
    "            nn.ELU(),\n",
    "            nn.Conv1d(int(kernel_count), int(kernel_count/2), 7, stride=1, padding=0),\n",
    "            nn.BatchNorm1d(int(kernel_count/2)),\n",
    "            nn.ELU(),\n",
    "            nn.Conv1d(int(kernel_count/2), int(kernel_count/4), 7, stride=2, padding=0),\n",
    "            nn.BatchNorm1d(int(kernel_count/4)),\n",
    "            nn.ELU(),\n",
    "            nn.Conv1d(int(kernel_count/4), 1, 7, stride=1, padding=0),\n",
    "            nn.BatchNorm1d(1),\n",
    "            nn.ELU()\n",
    "        )\n",
    "\n",
    "        self.decoder = nn.Sequential(\n",
    "            nn.ConvTranspose1d(1, int(kernel_count/4), 7, stride=1, padding=0),\n",
    "            nn.BatchNorm1d(int(kernel_count/4)),\n",
    "            nn.ELU(),\n",
    "            nn.ConvTranspose1d(int(kernel_count/4), int(kernel_count/2), 7, stride=2, padding=0, output_padding=1),\n",
    "            nn.BatchNorm1d(int(kernel_count/2)),\n",
    "            nn.ELU(),\n",
    "            nn.ConvTranspose1d(int(kernel_count/2), int(kernel_count), 7, stride=1, padding=0),\n",
    "            nn.BatchNorm1d(int(kernel_count)),\n",
    "            nn.ELU(),\n",
    "            nn.ConvTranspose1d(int(kernel_count), 1, 7, stride=1, padding=0), \n",
    "            nn.BatchNorm1d(1),\n",
    "            nn.ELU(),\n",
    "            nn.Conv1d(1, 1, 7, stride=1, padding='same'),\n",
    "            nn.Tanh()\n",
    "        )\n",
    "\n",
    "\n",
    "    def forward(self, x):\n",
    "        encoded = self.encoder(x)\n",
    "        decoded = self.decoder(encoded)\n",
    "        return decoded\n",
    "\n",
    "# Prints the structure of the Autoencoder\n",
    "torchinfo.summary(AE2(), (8, 1, window_size), col_names=('input_size', 'output_size'))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "fdc39d04",
   "metadata": {},
   "source": [
    "## Training\n",
    "\n",
    "The following cells define the training procedure. Training is carried out in a configurable number of epochs and with configurable batch sizes. Parameters that are relevant for the training procedure will be specified in a config, which will be passed to the training function via an argument. This includes the learning rate, weight decay, the batch size and the kernel count of the autoencoder models. MSE Loss will be used as a loss criterion, while the ADAM optimiser is used to optimise the parameters of the model. The training data will be split into a training (80 %) and validation set (20 %) and data loaders are used to load the data. The latter split the data into minibatches and also shuffle it for the usage in training and validation. In addition, the train function includes some features for ray tune hyperparameter tuning such as accessing the training data via a `data_id`. Finally, noise is added to the x samples in the training data before training and when desired."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c7816ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "Method used for training the autoencoders with a given config. \n",
    "\n",
    ":param config: The configuration for training includes the following parameters as a dictionary: \n",
    "                    - learning rate (lr, float)\n",
    "                    - weight decay (wd, float)\n",
    "                    - batch size (batch_size, int)\n",
    "                    - kernel count (kernel_count, int)\n",
    ":param train_dataset: The train dataset.\n",
    ":param model_num: The number of the autoencoder (1 or 2).\n",
    ":param directory: The relative directory inside the `model` directory to save the model and loss data in.\n",
    ":param use_tune: If true, this function is used with the ray tune hyperparameter tuning framework.\n",
    ":param data_id: Used by ray tune to access the data from the ray tune object store.\n",
    "'''\n",
    "def train(config, train_dataset, model_num, directory, use_tune=False, data_id=None):\n",
    "\n",
    "    save_dir = os.path.join(model_directory, directory)\n",
    "    if not os.path.exists(save_dir):\n",
    "        os.makedirs(save_dir)\n",
    "\n",
    "    # Select and create the correct model, send it to the correct device (cpu/cuda)\n",
    "    if model_num == 1:\n",
    "        model = AE1(kernel_count=config['kernel_count'])\n",
    "    elif model_num == 2:\n",
    "        model = AE2(kernel_count=config['kernel_count'])\n",
    "    else:\n",
    "        return\n",
    "    model = model.to(device)\n",
    "\n",
    "    # Use MSE Loss as the loss function and the ADAM optimiser for optimisation.\n",
    "    criterion = nn.MSELoss()\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=config['lr'], weight_decay=config['wd'])\n",
    "\n",
    "    # For ray tune. Get the train dataset from the ray tune object store with the given `data_id`.\n",
    "    if data_id:\n",
    "        train_dataset = get(data_id)\n",
    "\n",
    "    # Split the train dataset into a train and validation set.\n",
    "    val_size = int(val_prob * len(train_dataset))\n",
    "    train_size = len(train_dataset) - val_size\n",
    "    train_dataset, val_dataset = torch.utils.data.random_split(train_dataset, [train_size, val_size])\n",
    "\n",
    "    # Create DataLoaders for the train and validation dataset.\n",
    "    train_loader = DataLoader(dataset=train_dataset, batch_size=config['batch_size'], shuffle=True, num_workers=4, pin_memory=True)\n",
    "    val_loader = DataLoader(dataset=val_dataset, batch_size=config['batch_size'], shuffle=True, num_workers=4, pin_memory=True)\n",
    "\n",
    "    train_loss_values = []\n",
    "    val_loss_values = []\n",
    "    epoch_values = []\n",
    "\n",
    "    min_val_loss = np.inf\n",
    "\n",
    "    # Do training in several epochs\n",
    "    for epoch in range(num_epochs):\n",
    "        \n",
    "        train_loss = 0.0\n",
    "        for i, (inputs, targets) in enumerate(train_loader):\n",
    "            \n",
    "            # Send tensors to the desired device and reshape them to account for the input shape of the autoencoder.\n",
    "            inputs, targets = inputs.to(device, dtype=torch.float), targets.to(device, dtype=torch.float)\n",
    "            inputs, targets = inputs.view(-1, 1, window_size), targets.view(-1, 1, window_size)\n",
    "\n",
    "            # Calculate training loss and do backpropagation\n",
    "            optimizer.zero_grad()\n",
    "            output = model(inputs)\n",
    "            loss = criterion(output, targets)\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "            train_loss += loss.item()\n",
    "        \n",
    "        train_loss_values.append(train_loss)\n",
    "\n",
    "        val_loss = 0.0\n",
    "        for i, (inputs, targets) in enumerate(val_loader):\n",
    "\n",
    "            # Send tensors to the desired device and reshape them to account for the input shape of the autoencoder.\n",
    "            inputs, targets = inputs.to(device, dtype=torch.float), targets.to(device, dtype=torch.float)\n",
    "            inputs, targets = inputs.view(-1, 1, window_size), targets.view(-1, 1, window_size)\n",
    "\n",
    "            # Calculate validation loss\n",
    "            output = model(inputs)\n",
    "            loss = criterion(output, targets)\n",
    "            val_loss += loss.item()\n",
    "            \n",
    "        val_loss_values.append(val_loss)\n",
    "        epoch_values.append(epoch)\n",
    "\n",
    "        if use_tune:\n",
    "            # Pass the information of the current validation loss to ray tune.\n",
    "            tune.report(loss=(val_loss / len(val_loader)))\n",
    "        else:\n",
    "            # Print results from this epoch and save them in a CSV file\n",
    "            print(f'epoch {epoch + 1} / {num_epochs}, train loss = {(train_loss / len(train_loader)):.10f}, val loss = {(val_loss / len(val_loader)):.10f}')\n",
    "            pd.DataFrame({'epoch': epoch_values, 'train_loss': np.divide(train_loss_values, len(train_loader)), 'val_loss': np.divide(val_loss_values, len(val_loader))}).to_csv(os.path.join(save_dir, f'loss_{model_num}.csv')) \n",
    "            \n",
    "            # Save the model if validation loss got lower\n",
    "            if min_val_loss > val_loss:\n",
    "                print(f'\\tval loss decreased from {(min_val_loss/len(val_loader)):.10f} to {(val_loss/len(val_loader)):.10f}')\n",
    "                min_val_loss = val_loss\n",
    "                torch.save(model.state_dict(), os.path.join(save_dir, f'model_{model_num}.pth'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1549df4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Adds the correct noise levels for the training process.\n",
    "# train_dataset.dataset.add_noise(train_target_snrs, noise_model='WGN', wgn_prob=0.4, use_bandpass=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "354f4492",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.cuda.empty_cache()\n",
    "\n",
    "# Execute training \n",
    "# train({'lr': 0.01, 'wd': 0, 'batch_size': 8, 'kernel_count': 32}, train_dataset, 1, 'test')\n",
    "# train({'lr': 0.01, 'wd': 0, 'batch_size': 8, 'kernel_count': 256}, train_dataset, 2) # scg: 0.0015, ecg: 0.005"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "550d9c55",
   "metadata": {},
   "source": [
    "## Hyperparameter Tuning\n",
    "\n",
    "The following code cell implements the hyperparameter tuning which was developed using RayTune. For the optimisation of the hyperparameters, the `ASHAScheduler` is used and the task is to minimize the validation loss. The configuration defines the hyperparameters to be optimized. Here only the learning rate `lr`, the batch size `batch_size` and the kernel count `kernel_count` of the models will be taken into consideration. The weight decay remains at 0, as it was found that this parameter reduces the denoising performance of the autoencoders."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd2406ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "Method used for hyperparameter tuning of the machine learning models with a given config. \n",
    "\n",
    ":param dataset: The train dataset.\n",
    ":param config: The configuration for hyperparameter tuning as a dict. Includes the following parameters: \n",
    "                    - learning rate (lr, float)\n",
    "                    - weight decay (wd, float)\n",
    "                    - batch size (batch_size, int)\n",
    "                    - kernel count (kernel_count, int)\n",
    ":param num_samples: The number of different configurations of the config parameters to be tested.\n",
    ":param max_num_epochs: The maximum number of epochs that will be used for each sample during training.\n",
    ":param model_num: The number indicating the autoencoder (1 or 2).\n",
    "'''\n",
    "def hypertuning(dataset, config, num_samples, max_num_epochs, model_num):\n",
    "\n",
    "    ray.init()\n",
    "    \n",
    "    # As a scheduler, the ASHAScheduler is used with the task to minimise the (validation) loss.\n",
    "    scheduler = ASHAScheduler(metric='loss', mode='min', max_t=max_num_epochs, grace_period=4, reduction_factor=2)\n",
    "    # Reports the results to the console.\n",
    "    reporter = CLIReporter(metric_columns=['loss', 'training_iteration'])\n",
    "    # Put the dataset into the ray object store for faster and more efficient access.\n",
    "    data_id = put(dataset)\n",
    "    # Runs the scheduler with the train method. \n",
    "    result = tune.run(\n",
    "        partial(train, train_dataset=None, model_num=model_num, use_tune=True, data_id=data_id),\n",
    "        resources_per_trial={'cpu': 8, 'gpu': 1},\n",
    "        config=config,\n",
    "        num_samples=num_samples,\n",
    "        scheduler=scheduler,\n",
    "        progress_reporter=reporter\n",
    "    )\n",
    "\n",
    "    # Report the best config\n",
    "    best_trial = result.get_best_trial('loss', 'min', 'last')\n",
    "    print(\"Best trial config: {}\".format(best_trial.config))\n",
    "    print(\"Best trial final validation loss: {}\".format(best_trial.last_result[\"loss\"]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0cfe8d4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "config = {\n",
    "    'lr': tune.loguniform(1e-4, 1e-1),\n",
    "    'wd': 0,\n",
    "    'batch_size': tune.choice([8]),\n",
    "    'kernel_count': tune.choice([256])\n",
    "}\n",
    "\n",
    "# hypertuning(train_dataset, config, num_samples=16, max_num_epochs=32, model_num=2)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "21410168",
   "metadata": {},
   "source": [
    "## Testing and Evaluation\n",
    "\n",
    "The following code cells implement the testing and evaluation procedure. In the first cell, several evaluation metrics that are used in the subsequent functions are defined. Thereafter, some utility functions to plot and save the results are implemented. The main function, however, is the `test` function, which carries out the entire test and evaluation procedure. The function is called with a configuration as an argument, the specific contents of which are described later. This config defines the setup for the different experiments and sub-experiments. \n",
    "\n",
    "In the test function, all models specified in the config are tested against all evaluation metrics an all `test_target_snrs` with the test dataset and the results are saved in a results dict with the structure results[method][scope][evaluation_method]. The method can either be 'AE1' or 'AE2' for the autoencoder models trained for denoising, 'ECG' for the models trained for ECG/SCG transformation or 'WD' for Wavelet Denoising. The scope, on the other hand, can be either value from `test_target_snrs`, `all` for a collection of values over all SNR values, or `ECG Noise` or `ECG Predicted` for the last experiment when testing on the noisy or denoised/predicted SCG samples respectively. The name of the evaluation method coincides with the function names of the first code cell, with the SNR improvement being denoted as `snr_imp`. \n",
    "\n",
    "For the CPU time, only the prediction task itself is measured and any preprocessing or postprocessing tasks like reshaping the input tensors are neglected. However, if bandpass filters are used in the pre-processing, the time taken by them is added to the CPU time to allow a comparison in the first experiment. To measure the CPU time, the `process_time()` function is used which returns the sum of the system and user CPU time of the current process ([Source](https://docs.python.org/3/library/time.html#time.process_time)). Finally, the noise for testing is added to the test dataset before testing and only if it is desired, i.e. the parameter `noise` of the dataset is set to True."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d3bbad1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here, some evaluation metrics are implemented.\n",
    "\n",
    "# Mean-Squared Error\n",
    "def mse(signal_clean, signal_noisy):\n",
    "    return mean_squared_error(signal_clean, signal_noisy)\n",
    "\n",
    "# Mean-Absolute Error\n",
    "def mae(signal_clean, signal_noisy):\n",
    "    return mean_absolute_error(signal_clean, signal_noisy)\n",
    "\n",
    "# Root Mean-Squared Error\n",
    "def rmse(signal_clean, signal_noisy):\n",
    "    return np.sqrt(mse(signal_clean, signal_noisy))\n",
    "\n",
    "# Signal-to-noise ratio\n",
    "def snr(signal_clean, signal_noisy):\n",
    "    return 10 * np.log10(np.sum(signal_clean ** 2) / (np.sum((signal_noisy - signal_clean) ** 2)))\n",
    "\n",
    "# Root mean squared difference\n",
    "def prd(signal_clean, signal_noisy):\n",
    "    return np.sqrt(np.sum((signal_clean - signal_noisy) ** 2) / np.sum(signal_clean ** 2)) * 100\n",
    "\n",
    "# Cross-Correlation\n",
    "def xcorr(signal_clean, signal_noisy):\n",
    "    return np.mean(ccf(signal_clean, signal_noisy, adjusted=False))\n",
    "\n",
    "# Pearson Correlation Coefficient\n",
    "def pcorr(signal_clean, signal_noisy):\n",
    "    return np.corrcoef(signal_clean, signal_noisy)[1][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dfd42bab",
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "Utility function to plot a chart including multiple bar charts, each for a specific method in the results, with one bar for each denoted SNR.\n",
    "\n",
    ":param name_short: The short name of the evaluation method that coincides with the function name. The SNR improvement is denoted as `snr_imp`.\n",
    ":param results: The results dict as it was returned from the test method.\n",
    ":param labels: The labels for the bar charts on the y-axis as a list of strings. Must be the same length as the number of methods in the results.\n",
    ":param target_snrs_print: The SNRs that will be printed. This has to be a subset of `test_target_snrs`, `all` and [`ECG Noise`, `ECG Predicted`]\n",
    ":param xlabel: The label on the x-axis.\n",
    ":param ylabel: The label on the y-axis.\n",
    ":param title: The title of the chart.\n",
    ":param directory: The directory where the chart will be saved as a png.\n",
    ":param printerror: Whether an error bar should be printed on top of the bar charts to represent the standard deviation.\n",
    ":param subplot: Whether the plot should be added to a subplot passed to the function by this argument. None if it should be a separate plot.\n",
    "'''\n",
    "def plot_result_bar(name_short, results, labels, target_snrs_print, xlabel, ylabel, title, directory, printerror=False, subplot=None):\n",
    "    \n",
    "    # Creates lists for the standard deviations and means of the respective data to pass it directly to the functions that draw the charts.\n",
    "    data_mean = [[results[method][scope][name_short]['mean'] for method in results.keys()] for scope in target_snrs_print]\n",
    "    data_std = [[results[method][scope][name_short]['sd'] for method in results.keys()] for scope in target_snrs_print]\n",
    "\n",
    "    if subplot:\n",
    "        # If a subplot was passed to the function, add the plot to it.\n",
    "        subplot.set_xticks(range(len(data_mean[0])), labels)\n",
    "        subplot.set_xlabel(xlabel)\n",
    "        subplot.set_ylabel(ylabel)\n",
    "        subplot.set_title(title)\n",
    "    else:\n",
    "        # Create a normal plot otherwise.\n",
    "        figsize = (6.4, 4.8) if name_short == 'snr_imp' else (6.4, 2.4)\n",
    "        plt.figure(figsize=figsize)\n",
    "        plt.xticks(range(len(data_mean[0])), labels)\n",
    "        plt.xlabel(xlabel)\n",
    "        plt.ylabel(ylabel)\n",
    "        plt.title(title)\n",
    "    \n",
    "    # The width of the bars in the plots\n",
    "    width = 0.12 if len(target_snrs_print) > 2 else 0.3 if len(target_snrs_print) == 2 else 0.6\n",
    "    # The colors used for the bars\n",
    "    color_map = ['#004a99', '#0065b0', '#117fc6', '#2c9bda', '#49b6ed', '#67d2ff'] if len(target_snrs_print) == 6 else ['#0065b0', '#2c9bda'] if len(target_snrs_print) == 2 else ['#0065b0']\n",
    "    \n",
    "    for i in range(len(target_snrs_print)):\n",
    "\n",
    "        x = np.arange(len(data_mean[i])) + i * width\n",
    "        y = data_mean[i]\n",
    "        e = data_std[i]\n",
    "\n",
    "        if subplot:\n",
    "            subplot.bar(x, y, width=width, label=target_snrs_print[i], color=color_map[i])\n",
    "            if printerror:\n",
    "                subplot.errorbar(x, y, e, linestyle='None', fmt ='o', capsize=5, ecolor='black', markerfacecolor='black', markeredgecolor='black')\n",
    "        else: \n",
    "            plt.bar(x, y, width=width, label=target_snrs_print[i], color=color_map[i])\n",
    "            if printerror:\n",
    "                plt.errorbar(x, y, e, linestyle='None', fmt ='o', capsize=5, ecolor='black', markerfacecolor='black', markeredgecolor='black')\n",
    "\n",
    "        loc = 'lower right' if name_short in ['rmse', 'pcorr'] else 'best'\n",
    "        plt.legend(ncol=2, loc=loc)\n",
    "    \n",
    "    if not subplot:\n",
    "        plt.plot()\n",
    "        plt.savefig(os.path.join(directory, name_short + '.png'), format='png', bbox_inches='tight')\n",
    "        plt.show()\n",
    "\n",
    "'''\n",
    "Utility function to save the results dict as returned by the test method as a json and to plot the results. \n",
    "\n",
    ":param directory: The directory where the results should be saved.\n",
    ":param config: The config dict from the test method.\n",
    ":param results: The results dict that was returned from the test method.\n",
    "'''\n",
    "def save_results(directory, config, results=None):\n",
    "\n",
    "    if not os.path.exists(directory):\n",
    "        os.makedirs(directory)\n",
    "\n",
    "    filename = os.path.join(directory, 'results.json')\n",
    "    if results and not os.path.exists(filename):\n",
    "        # Create a dictionary from the result data with calculated mean and standard deviation and save it into a json file\n",
    "        results_print = {method: {scope: {evaluation: {'mean': float(np.mean(values)), 'sd': float(np.std(values))} for evaluation, values in eval_method.items()} for scope, eval_method in method_value.items()} for method, method_value in results.items()}\n",
    "        json.dump(results_print, open(filename, 'w+'), indent=4)\n",
    "    elif directory == os.path.join('.', 'results', '3', 'complete'):\n",
    "        # Creates the results.json for the last experiment that was carried out in three different stages (on three datasets)\n",
    "        results_1 = json.load(open(os.path.join('.', 'results', '3', 'improvement_1', 'results.json')))\n",
    "        results_2 = json.load(open(os.path.join('.', 'results', '3', 'improvement_2', 'results.json')))\n",
    "        results_3 = json.load(open(os.path.join('.', 'results', '3', 'improvement_3', 'results.json')))\n",
    "\n",
    "        results_print = {'WGN -1': results_1['AE2'], 'WGN 7': results_2['AE2'], 'Motion': results_3['AE2']}\n",
    "        json.dump(results_print, open(filename, 'w+'), indent=4)\n",
    "    else:\n",
    "        # Load the results if they are already calculated\n",
    "        results_print = json.load(open(filename))\n",
    "\n",
    "\n",
    "    # Plot the different bar charts for the evaluation metrics\n",
    "    plot_result_bar('snr_imp', results_print, config['labels'], config['target_snrs_print'], config['xlabel'], 'SNR Improvement (dB)', 'Mean SNR Improvement', directory, printerror=False)\n",
    "    plot_result_bar('mse', results_print, config['labels'], config['target_snrs_print'], config['xlabel'], 'MSE', 'Mean MSE', directory, printerror=False)\n",
    "    plot_result_bar('rmse', results_print, config['labels'], config['target_snrs_print'], config['xlabel'], 'RMSE', 'Mean RMSE', directory, printerror=False)\n",
    "    plot_result_bar('mae', results_print, config['labels'], config['target_snrs_print'], config['xlabel'], 'MAE', 'Mean MAE', directory, printerror=False)\n",
    "    plot_result_bar('xcorr', results_print, config['labels'], config['target_snrs_print'], config['xlabel'], 'xcorr', 'Mean xcorr', directory, printerror=False)\n",
    "    plot_result_bar('pcorr', results_print, config['labels'], config['target_snrs_print'], config['xlabel'], 'pcorr', 'Mean pcorr', directory, printerror=False)\n",
    "    plot_result_bar('prd', results_print, config['labels'], config['target_snrs_print'], config['xlabel'], 'PRD (%)', 'Mean PRD', directory, printerror=False)\n",
    "\n",
    "\n",
    "    # Plot that contains the SNR improvement, RMSE and pcorr as subplots\n",
    "    plt.figure(figsize=(16, 6))\n",
    "    \n",
    "    ax1 = plt.subplot2grid((2, 8), (0, 0), colspan=4, rowspan=2)\n",
    "    plot_result_bar('snr_imp', results_print, config['labels'], config['target_snrs_print'], config['xlabel'], 'SNR Improvement (dB)', 'Mean SNR Improvement', directory, printerror=False, subplot=ax1)\n",
    "    \n",
    "    ax2 = plt.subplot2grid((2, 8), (0, 4), colspan=4, rowspan=1)\n",
    "    ax2.tick_params(labelbottom=False)\n",
    "    plot_result_bar('rmse', results_print, config['labels'], config['target_snrs_print'], config['xlabel'], 'RMSE', 'Mean RMSE', directory, printerror=False, subplot=ax2)\n",
    "    ax2.set(xlabel=None)\n",
    "    \n",
    "    ax3 = plt.subplot2grid((2, 8), (1, 4), colspan=4, rowspan=1)\n",
    "    plot_result_bar('pcorr', results_print, config['labels'], config['target_snrs_print'], config['xlabel'], 'pcorr', 'Mean pcorr', directory, printerror=False, subplot=ax3)\n",
    "    \n",
    "    plt.tight_layout()\n",
    "    plt.plot()\n",
    "    plt.savefig(os.path.join(directory, 'complete.png'), format='png', bbox_inches='tight')\n",
    "\n",
    "\n",
    "    # Plot for the CPU consumption time\n",
    "    plt.figure(figsize=(6.4, 2.4))\n",
    "    data = [results_print[method]['all']['cpu_time']['mean'] * 1000 for method in results_print.keys()]\n",
    "\n",
    "    if config['labels'][2] == 'Train/Test Data + Results':\n",
    "        config['labels'][2] = 'Train/Test Data +\\n Results'\n",
    "        \n",
    "    plt.xticks(range(len(data)), config['labels'])\n",
    "    plt.xlabel(config['xlabel'])\n",
    "    plt.ylabel('CPU Time (ms)')\n",
    "    plt.title('Mean CPU Time')\n",
    "\n",
    "    plt.bar(np.arange(len(data)), data, color='#117fc6')\n",
    "    \n",
    "    plt.savefig(os.path.join(directory, 'cpu.png'), format='png', bbox_inches='tight')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ab33828",
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "Method for testing the denoising methods in relation to a specific configuration for each (sub-)experiment.\n",
    "\n",
    ":param config: The configuration for testing as a dict. Includes the following parameters: \n",
    "                    - methods           (list of strings)   The used methods. The value may be 'AE1', 'AE2', 'ECG' or 'WD' for the first or second autoencoder, the second autoencoder trained for SCG/ECG transformation and wavelet denoising respectively\n",
    "                    - models            (list of strings)   One value per method. The relative path to the saved model in the models directory that matches the corresponding method at the same location in the list\n",
    "                    - labels            (list of strings)   One value per method. The labels for the different methods in the diagrams and results\n",
    "                    - target_snrs_print (list of ints)      The SNRs that should be printed in the diagrams\n",
    "                    - bandpass_results  (list of bools)     One value per method. True, if a bandpass filter should be applied to the results of the respective method, False otherwise.\n",
    "                    - bandpass_inputs   (list of bools)     One value per method. True, if a bandpass filter should be applied to the possibly noisy input data of the respective method, False otherwise.\n",
    "                    - ecg               (bool)              If set to True, the performance of the denoising model will be tested on the transformation from SCG to ECG signals\n",
    "                    - ecg_model         (string)            The path to the transformation model for SCG to ECG if `ecg` is set to True\n",
    "                    - xlabel            (string)            The label for the x-axis in the diagram\n",
    "'''\n",
    "\n",
    "def test(config):\n",
    "    \n",
    "    with torch.no_grad():\n",
    "\n",
    "        # A dict to save the results in\n",
    "        results = {}\n",
    "\n",
    "        if config['ecg']:\n",
    "            # Loads the model for the transformation from scg to ecg\n",
    "            model_ecg = AE2(kernel_count=256).to(device)\n",
    "            model_ecg.load_state_dict(torch.load(os.path.join(model_directory, config['ecg_model']), map_location=device))\n",
    "            model_ecg.eval()\n",
    "\n",
    "        for i, method in enumerate(config['methods']):\n",
    "\n",
    "            label = config['labels'][i]\n",
    "\n",
    "            # Initialise the current machine learning model\n",
    "            model = None\n",
    "\n",
    "            if method == 'AE1':\n",
    "                model = AE1(kernel_count=128).to(device)\n",
    "            elif method == 'AE2':\n",
    "                model = AE2(kernel_count=128).to(device)\n",
    "            elif method == 'ECG':\n",
    "                model = AE2(kernel_count=256).to(device)\n",
    "            \n",
    "            if model:\n",
    "                model.load_state_dict(torch.load(os.path.join(model_directory, config['models'][i]), map_location=device))\n",
    "                model.eval()\n",
    "\n",
    "            # Initialise the results dict for the different parameters to be tested\n",
    "            results[label] = {}\n",
    "            for target in test_target_snrs + ['all', 'ECG Noise', 'ECG Predicted']:\n",
    "                target = str(target)\n",
    "                results[label][target] = {\n",
    "                    'result': [],\n",
    "                    'snr_imp': [],\n",
    "                    'mse': [],\n",
    "                    'rmse': [],\n",
    "                    'mae': [],\n",
    "                    'xcorr': [],\n",
    "                    'pcorr': [],\n",
    "                    'prd': [],\n",
    "                    'cpu_time': []\n",
    "                }\n",
    "            \n",
    "            # Calculate the results for each sample in the test dataset\n",
    "            for j, (inputs, targets) in enumerate(test_dataset):\n",
    "                \n",
    "                cpu_time = 0\n",
    "                # Inputs for the prediction/denoising task\n",
    "                inputs_snr = inputs.detach().clone().numpy()\n",
    "\n",
    "                if config['bandpass_inputs'][i]:\n",
    "                    # Apply the bandpass filter to the noisy input data if required\n",
    "                    inputs = inputs.numpy()\n",
    "                    cpu_time_start = time.process_time()\n",
    "                    inputs = bandpass(inputs, order=4, fs=100.0, lowcut=5.0, highcut=30.0)\n",
    "                    cpu_time_stop = time.process_time()\n",
    "                    cpu_time += cpu_time_stop - cpu_time_start\n",
    "                    inputs = torch.from_numpy(inputs.copy())\n",
    "\n",
    "                # Test the corresponding method\n",
    "                if method in ['AE1', 'AE2', 'ECG']:\n",
    "                    inputs = inputs.reshape(-1, 1, window_size).to(device, dtype=torch.float)\n",
    "\n",
    "                    # Record the cpu time as the process time\n",
    "                    cpu_time_start = time.process_time()\n",
    "                    predicted = model(inputs)\n",
    "                    cpu_time_stop = time.process_time()\n",
    "                    cpu_time += cpu_time_stop - cpu_time_start\n",
    "\n",
    "                    predicted = predicted.to('cpu').reshape(window_size).numpy()\n",
    "                    inputs = inputs.to('cpu').reshape(window_size).numpy()\n",
    "                elif method == 'WD':\n",
    "                    inputs = inputs.to('cpu').reshape(window_size).numpy()\n",
    "\n",
    "                    # Record the cpu time as the process time\n",
    "                    cpu_time_start = time.process_time()\n",
    "                    predicted = denoise_wavelet(inputs, method='BayesShrink', mode='soft', wavelet='sym8', rescale_sigma='True')\n",
    "                    cpu_time_stop = time.process_time()\n",
    "                    cpu_time += cpu_time_stop - cpu_time_start\n",
    "                    # Denormalise, because the target distribution of wavelet denoising is not equivalent to the denormalised values as with the autoencoders\n",
    "                    predicted = denormalise(predicted, test_dataset.dataset.min_val, test_dataset.dataset.max_val)\n",
    "\n",
    "                targets = targets.numpy()\n",
    "\n",
    "                if config['bandpass_results'][i]:\n",
    "                    # Apply the bandpass filter to the results if required\n",
    "                    cpu_time_start = time.process_time()\n",
    "                    predicted = bandpass(predicted, order=4, fs=100.0, lowcut=5.0, highcut=30.0)\n",
    "                    cpu_time_stop = time.process_time()\n",
    "                    cpu_time += cpu_time_stop - cpu_time_start\n",
    "\n",
    "                # Last experiment: Test the transformation autoencoder on the predicted/denoised and noisy scg signal\n",
    "                if config['ecg']:\n",
    "                    predicted = torch.from_numpy(predicted).reshape(-1, 1, window_size).to(device, dtype=torch.float) # denoised\n",
    "                    inputs = torch.from_numpy(inputs).reshape(-1, 1, window_size).to(device, dtype=torch.float) # noisy\n",
    "                    \n",
    "                    # Get the transformed signal from the noisy and predicted/denoisied scg signal\n",
    "                    predicted_ecg_inputs = model_ecg(inputs).to('cpu').reshape(window_size).numpy() # prediction on noisy data\n",
    "                    predicted_ecg_pred = model_ecg(predicted).to('cpu').reshape(window_size).numpy() # prediction on denoised data\n",
    "\n",
    "                    # Calculate the RMSE and pcorr for the ecg/scg transformation for the noisy data and predicted/denoised scg data\n",
    "                    results[label]['ECG Noise']['rmse'].append([rmse(targets, predicted_ecg_inputs)])\n",
    "                    results[label]['ECG Noise']['pcorr'].append([pcorr(targets, predicted_ecg_inputs)]) \n",
    "\n",
    "                    results[label]['ECG Predicted']['rmse'].append([rmse(targets, predicted_ecg_pred)])\n",
    "                    results[label]['ECG Predicted']['pcorr'].append([pcorr(targets, predicted_ecg_pred)]) \n",
    "                \n",
    "                # Test for the remaining experiments\n",
    "                else:\n",
    "                    # Denormalise the original input values for the calculation of the SNR as it depends on absolute values.\n",
    "                    if test_dataset.dataset.noise:\n",
    "                        inputs_snr = denormalise(inputs_snr, test_dataset.dataset.min_val, test_dataset.dataset.max_val)\n",
    "\n",
    "                    # Calculate the performance measures for each SNR on the test dataset. The SNR values are selected with a margin of 0.8.\n",
    "                    snr_data = snr(targets, inputs_snr)\n",
    "                    for target in test_target_snrs + ['all']:\n",
    "                        if (str(target) == 'all') or (snr_data > (target - 0.8) and snr_data < (target + 0.8)):\n",
    "                            j += 1\n",
    "                            target = str(target)\n",
    "                            results[label][target]['result'].append(predicted)\n",
    "                            results[label][target]['snr_imp'].append([snr(targets, predicted) - snr_data])\n",
    "                            results[label][target]['mse'].append([mse(targets, predicted)])\n",
    "                            results[label][target]['rmse'].append([rmse(targets, predicted)])\n",
    "                            results[label][target]['mae'].append([mae(targets, predicted)])\n",
    "                            results[label][target]['xcorr'].append([xcorr(targets, predicted)])\n",
    "                            results[label][target]['pcorr'].append([pcorr(targets, predicted)])\n",
    "                            results[label][target]['prd'].append([prd(targets, predicted)])\n",
    "                            results[label][target]['cpu_time'].append([cpu_time])\n",
    "            \n",
    "            if config['ecg']:\n",
    "                print(f'Label: {label}')\n",
    "                print(f'Values: {j}')\n",
    "                print(len(results[label][\"ECG Predicted\"][\"result\"]))\n",
    "                print(f'RMSE {str(np.mean(results[label][\"ECG Noise\"][\"rmse\"]))}/{str(np.mean(results[label][\"ECG Predicted\"][\"rmse\"]))}')\n",
    "                print(f'PCORR {str(np.mean(results[label][\"ECG Noise\"][\"pcorr\"]))}/{str(np.mean(results[label][\"ECG Predicted\"][\"pcorr\"]))}')\n",
    "            else:\n",
    "                value = 'all'\n",
    "                print(f'Label: {label}')\n",
    "                print(f'Values: {j}')\n",
    "                print(len(results[label][value][\"result\"]))\n",
    "                print(f'SNR IMP {str(np.mean(results[label][value][\"snr_imp\"]))}')\n",
    "                print(f'RMSE {str(np.mean(results[label][value][\"rmse\"]))}')\n",
    "                print(f'PCORR {str(np.mean(results[label][value][\"pcorr\"]))}')\n",
    "                print(f'CPU TIME {str(np.mean(results[label][value][\"cpu_time\"]))}\\n')\n",
    "\n",
    "        return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3266b49e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Testing is done with a different set of target SNRs. Those are set here.\n",
    "test_dataset.dataset.add_noise(test_target_snrs, noise_model='WGN', wgn_prob=0.4, use_bandpass=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8fbe20c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The config describes, which parameters should be tested.\n",
    "\n",
    "# Config for the first experiment\n",
    "config0 = {\n",
    "    'methods': ['AE2', 'AE2', 'AE2', 'AE2'], #, 'WD']\n",
    "    'models': ['wgn/model_2_wgn3_bandpass.pth', 'wgn/model_2_wgn3.pth', 'wgn/model_2_wgn3_bandpass.pth', 'wgn/model_2_wgn3.pth'],\n",
    "    'labels': ['Train/Test Data', 'Results', 'Train/Test Data + Results', 'No Usage'],\n",
    "    'target_snrs_print': ['-1', '0', '2', '5', '7', 'all'], # [-1, 0, 2, 5, 7]\n",
    "    'bandpass_results': [False, True, True, False],\n",
    "    'bandpass_inputs': [True, False, True, False],\n",
    "    'ecg': False,\n",
    "    'ecg_model': None,\n",
    "    'xlabel': 'Bandpass Filter Usage'\n",
    "}\n",
    "\n",
    "# Config for the second experiment\n",
    "config1 = {\n",
    "    'methods': ['AE2', 'AE2', 'AE2', 'AE2', 'AE2'], #, 'WD']\n",
    "    'models': ['wgn/model_2_clean.pth', 'wgn/model_2_wgn1.pth', 'wgn/model_2_wgn2.pth', 'wgn/model_2_wgn3.pth', 'wgn/model_2_wgn4.pth'],\n",
    "    'labels': ['None', 'WGN1', 'WGN2', 'WGN3', 'WGN4'],\n",
    "    'target_snrs_print': ['-1', '0', '2', '5', '7', 'all'], # [-1, 0, 2, 5, 7]\n",
    "    'bandpass_results': [False, False, False, False, False],\n",
    "    'bandpass_inputs': [False, False, False, False, False],\n",
    "    'ecg': False,\n",
    "    'ecg_model': None,\n",
    "    'xlabel': 'Noise Distribution'\n",
    "}\n",
    "\n",
    "# Config for the third experiment regarding the WGN noise model\n",
    "config2_1 = {\n",
    "    'methods': ['AE1', 'AE2', 'WD'],\n",
    "    'models': ['wgn/model_1_wgn3.pth', 'wgn/model_2_wgn3.pth', ''],\n",
    "    'labels': ['AE1', 'AE2', 'WD'],\n",
    "    'target_snrs_print': ['-1', '0', '2', '5', '7', 'all'], # [-1, 0, 2, 5, 7]\n",
    "    'bandpass_results': [False, False, False],\n",
    "    'bandpass_inputs': [False, False, False],\n",
    "    'ecg': False,\n",
    "    'ecg_model': None,\n",
    "    'xlabel': 'Denoising Method'\n",
    "}\n",
    "\n",
    "# Config for the third experiment regarding the motion-based noise model\n",
    "config2_2 = {\n",
    "    'methods': ['AE1', 'AE2', 'WD'],\n",
    "    'models': ['motion/model_1_wgn3.pth', 'motion/model_2_wgn3.pth', ''],\n",
    "    'labels': ['AE1', 'AE2', 'WD'],\n",
    "    'target_snrs_print': ['all'], # [-1, 0, 2, 5, 7]\n",
    "    'bandpass_results': [False, False, False],\n",
    "    'bandpass_inputs': [False, False, False],\n",
    "    'ecg': False,\n",
    "    'ecg_model': None,\n",
    "    'xlabel': 'Denoising Method'\n",
    "}\n",
    "\n",
    "# Config for the third experiment regarding verification experiment\n",
    "config2_3 = {\n",
    "    'methods': ['AE2', 'AE2'],\n",
    "    'models': ['wgn/model_2_wgn3_cebs_full.pth', 'wgn/model_2_wgn3_cebs_omitted.pth'],\n",
    "    'labels': ['CEBS Dataset', 'CEBS Dataset Without Subject 1 and 2'],\n",
    "    'target_snrs_print': ['-1', '0', '2', '5', '7', 'all'], # [-1, 0, 2, 5, 7]\n",
    "    'bandpass_results': [False, False, False],\n",
    "    'bandpass_inputs': [False, False, False],\n",
    "    'ecg': False,\n",
    "    'ecg_model': None,\n",
    "    'xlabel': 'Training Scope'\n",
    "}\n",
    "\n",
    "# Config for the fourth experiment regarding performance of the SCG/ECG transform\n",
    "config3_1 = {\n",
    "    'methods': ['ECG'],\n",
    "    'models': ['ecg/model_ecg_both_norm.pth'],\n",
    "    'labels': ['ECG'],\n",
    "    'target_snrs_print': ['all'], # [-1, 0, 2, 5, 7]\n",
    "    'bandpass_results': [False, False, False, False],\n",
    "    'bandpass_inputs': [False, False, False, False],\n",
    "    'ecg': False,\n",
    "    'ecg_model': None,\n",
    "    'xlabel': 'Denoising Method'\n",
    "}\n",
    "\n",
    "# Config for the fourth experiment regarding the SCG/ECG transformation performance improvement\n",
    "config3_2 = {\n",
    "    'methods': ['AE2'],\n",
    "    'models': ['motion/model_2_wgn3.pth'],\n",
    "    'labels': ['AE2'],\n",
    "    'target_snrs_print': ['ECG Noise', 'ECG Predicted'], # [-1, 0, 0.5, 1, 3, 5, 8]\n",
    "    'bandpass_results': [False, False, False, False],\n",
    "    'bandpass_inputs': [False, False, False, False],\n",
    "    'ecg': True,\n",
    "    'ecg_model': 'ecg/model_ecg_scg_norm.pth',\n",
    "    'xlabel': 'SCG Source'\n",
    "}\n",
    "\n",
    "# Config for the fourth experiment to compile the results from the other training runs with config_3_2. \n",
    "# This config should only be used with the save_results function, but not with the test function itself.\n",
    "config3_complete = {\n",
    "    'methods': ['AE2'],\n",
    "    'models': ['motion/model_2_wgn3.pth'],\n",
    "    'labels': ['WGN -1', 'WGN 7', 'Motion'],\n",
    "    'target_snrs_print': ['ECG Noise', 'ECG Predicted'], # [-1, 0, 0.5, 1, 3, 5, 8]\n",
    "    'bandpass_results': [False, False, False, False],\n",
    "    'bandpass_inputs': [False, False, False, False],\n",
    "    'ecg': True,\n",
    "    'ecg_model': 'ecg/model_ecg.pth',\n",
    "    'xlabel': 'Noise Model'\n",
    "}\n",
    "\n",
    "results_eval = test(config2_1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a22b77b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save and print the results\n",
    "save_results('./results/3/improvement_4', config2_3, results=results_eval)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f6c2eda4",
   "metadata": {},
   "source": [
    "## Utility Functions for Evaluation and Visualisation\n",
    "\n",
    "The following cells contain code to generate visualisations or to extract important information from the results for different parts of the bachelor thesis. This includes visualisations of the denoising performance and the noise models, methods to generate tables of the results, and measures to compare the different denoising methods."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10269bec",
   "metadata": {},
   "outputs": [],
   "source": [
    "# VISUALISATIONS OF THE LEARNING BEHAVIOR\n",
    "\n",
    "# Visualises the the validation and train loss of the given loss_*.csv in a single graph\n",
    "\n",
    "# filename = 'models/ecg/loss_ecg_both_norm.csv'\n",
    "# filename = 'models/motion/loss_1_10_motion.csv'\n",
    "filename = 'models/wgn/loss_2_wgn3_bandpass.csv'\n",
    "\n",
    "directory = os.path.join('results', 'training')\n",
    "if not os.path.exists(directory):\n",
    "    os.makedirs(directory)\n",
    "\n",
    "df = pd.read_csv(filename)\n",
    "\n",
    "plt.plot(df['epoch'], df['val_loss'], color='#087ffb', label='validation loss')\n",
    "plt.plot(df['epoch'], df['train_loss'], color='#ea002e', label='train loss')\n",
    "\n",
    "plt.title('Validation and Train Loss of AE2 \\nwith Added WGN3 Noise Filtered with a Bandpass Filter')\n",
    "plt.xlabel('Epoch')\n",
    "plt.ylabel('Loss')\n",
    "plt.legend()\n",
    "\n",
    "plt.plot()\n",
    "plt.savefig(os.path.join(directory, 'scg_2_wgn3_bandpass.png'), format='png', bbox_inches='tight')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc3d34b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# VISUALISATIONS FOR THE DENOISING PERFORMANCE AND THE NOISE MODEL\n",
    "\n",
    "# Visualises a clean and noisy SCG signal in one graph and the clean and denoised/transformed SCG/ECG signal in another graph. \n",
    "# The signal will be selected randomly and the denoising method can be configured.\n",
    "\n",
    "# Whether the SCG/ECG transformation performance with possibly noisy SCG signals should be visualised\n",
    "ecg = True\n",
    "\n",
    "# True, if wavelet denoising should be visualised \n",
    "wd = False\n",
    "\n",
    "directory = os.path.join('results', 'denoising_viz')\n",
    "if not os.path.exists(directory):\n",
    "    os.makedirs(directory)\n",
    "\n",
    "# The model for denoising. Either AE1 or AE2\n",
    "\n",
    "# model = AE1(kernel_count=256).to(device)\n",
    "# model.load_state_dict(torch.load(os.path.join(model_directory, 'wgn/m2_128_20.pth'), map_location=device))\n",
    "\n",
    "model = AE2(kernel_count=128).to(device)\n",
    "model.load_state_dict(torch.load(os.path.join(model_directory, 'motion/model_2_wgn3.pth'), map_location=device))\n",
    "\n",
    "# AE2 as the model for SCG/ECG transformation\n",
    "model_ecg = AE2(kernel_count=256).to(device)\n",
    "model_ecg.load_state_dict(torch.load(os.path.join(model_directory, 'ecg/model_ecg_scg_norm.pth'), map_location=device))\n",
    "\n",
    "model.eval()\n",
    "model_ecg.eval()\n",
    "\n",
    "with torch.no_grad():\n",
    "    \n",
    "    # Randomly select a sample from the test dataset.\n",
    "    i = random.choice(range(0, len(test_dataset)))\n",
    "    # i = 13498 # WGN\n",
    "    # i = 11652 # Motion\n",
    "    i = 20827 # ECG\n",
    "    print(f'Sample index: {i}')\n",
    "\n",
    "    noise, clean = test_dataset[i]\n",
    "\n",
    "    if ecg:\n",
    "        # Only add noise here so that the SNR can be calculated\n",
    "        clean_scg = noise.numpy()\n",
    "        noise = test_dataset.dataset._add_motion(clean_scg)\n",
    "        snr_ecg = snr(clean_scg, noise)\n",
    "        noise = torch.from_numpy(noise.values)\n",
    "\n",
    "    # Denoise with either WD or the respective AE and predict the ECG from the results\n",
    "    if wd:\n",
    "        predicted = denoise_wavelet(noise.numpy(), method='BayesShrink', mode='soft', wavelet_levels=1, wavelet='sym8', rescale_sigma='True')\n",
    "        predicted = denormalise(predicted, test_dataset.dataset.min_val, test_dataset.dataset.max_val)\n",
    "        predicted = torch.from_numpy(predicted).reshape(-1, 1, window_size).to(device, dtype=torch.float)\n",
    "\n",
    "    clean, noise = clean.reshape(-1, 1, window_size).to(device, dtype=torch.float), noise.reshape(-1, 1, window_size).to(device, dtype=torch.float)\n",
    "\n",
    "    if not wd:\n",
    "        predicted = model(noise)\n",
    "    predicted_ecg = model_ecg(predicted)\n",
    "\n",
    "    predicted = predicted.to('cpu').reshape(window_size).numpy()\n",
    "    predicted_ecg = predicted_ecg.to('cpu').reshape(window_size).numpy()\n",
    "    clean = clean.to('cpu').reshape(window_size).numpy()\n",
    "    noise = noise.to('cpu').reshape(window_size).numpy()\n",
    "\n",
    "    # Denormalise the noisy data to calculate the SNR and use it and the denoised SCG for visualisation when not currently in the ECG mode. \n",
    "    # Use the normalised noise data and the predicted ECG otherwise.\n",
    "    noise_snr = noise if ecg else denormalise(noise, test_dataset.dataset.min_val, test_dataset.dataset.max_val)\n",
    "    predicted = predicted_ecg if ecg else predicted\n",
    "\n",
    "    print(f'RMSE \\nClean/Noise:\\t\\t{rmse(clean, noise)}\\nClean/Predicted:\\t{rmse(clean, predicted)}\\n\\n')\n",
    "    print(f'SNR \\nClean/Noise:\\t\\t{snr_ecg if ecg else snr(clean, noise_snr)}\\nClean/Predicted:\\t{snr(clean, predicted)}\\nClean/BP:\\t\\t{snr(clean, bandpass(predicted, order=4, fs=100.0, lowcut=5.0, highcut=30.0))}\\n\\n')\n",
    "    print(f'PCORR \\nClean/Noise:\\t\\t{pcorr(clean, noise)}\\nClean/Predicted:\\t{pcorr(clean, predicted)}\\nClean/BP:\\t\\t{pcorr(clean, bandpass(predicted, order=4, fs=100.0, lowcut=5.0, highcut=30.0))}\\n\\n')\n",
    "    \n",
    "    # The sample range that should be visualised\n",
    "    sample_from = 0\n",
    "    time_to = 512\n",
    "\n",
    "    # Clean/Noise Graph\n",
    "    plt.figure(figsize=(14, 4))\n",
    "\n",
    "    x = np.array(range(time_to)[sample_from:time_to]) / sampling_rate\n",
    "    plt.plot(x, clean_scg[sample_from:time_to] if ecg else clean[sample_from:time_to], color='#087ffb', label='clean')\n",
    "    plt.plot(x, noise_snr[sample_from:time_to], color='#ea002e', label='noisy')\n",
    "    plt.xticks(x[::25])\n",
    "    \n",
    "    plt.title('Signal with Added Motion Noise')\n",
    "    plt.xlabel('Time (s)')\n",
    "    plt.ylabel(r'Acceleration')\n",
    "    plt.ylim(bottom=-1, top=1)\n",
    "    plt.legend()\n",
    "\n",
    "    plt.plot()\n",
    "    plt.savefig(os.path.join(directory, 'noise_ecg.png'), format='png', bbox_inches='tight')\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "    # Clean/Denoised Graph\n",
    "    plt.figure(figsize=(14, 4))\n",
    "\n",
    "    x = np.array(range(time_to)[sample_from:time_to]) / sampling_rate\n",
    "    plt.plot(x, clean[sample_from:time_to], color='#087ffb', label='clean')\n",
    "    plt.plot(x, predicted[sample_from:time_to], color='#ea002e', label='denoised and transformed')\n",
    "    plt.xticks(x[::25])\n",
    "    \n",
    "    plt.title('Signal that has been Denoised with AE2 and Transformed into an ECG')\n",
    "    plt.xlabel('Time (s)')\n",
    "    plt.ylabel(r'Amplitude')\n",
    "    plt.ylim(bottom=-1, top=1)\n",
    "    plt.legend()\n",
    "\n",
    "    plt.plot()\n",
    "    plt.savefig(os.path.join(directory, 'denoising_ecg.png'), format='png', bbox_inches='tight')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65e47b1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# COMPARISON OF DIFFERENT DENOISING METHODS\n",
    "\n",
    "# Gives the  percentage change of the two specified labels (`label_1` and `label_2`) from the config of the respective `results.json` for specific scopes and a specific metric (mean or sd)\n",
    "\n",
    "# The filename for the respective results\n",
    "filename = 'results/2/motion/results.json'\n",
    "\n",
    "label_1 = 'WD' # The first label from the results\n",
    "label_2 = 'AE2' # The second label from the results\n",
    "scopes = ['all'] # The scope of SNRs\n",
    "metric = 'mean' # The metric to be shown (mean or sd)\n",
    "# 'labels': ['Training Data', 'Results', 'Training Data + Results', 'No Usage'],\n",
    "\n",
    "with open(filename) as json_file:\n",
    "    results = json.load(json_file)\n",
    "\n",
    "    for scope in scopes:\n",
    "        print(f'Scope: {scope}')\n",
    "        print('\\033[1m{:^8} {:^6} {:^6} {:^8}\\033[0m'.format('Evaluation', label_1, label_2, 'Percentage Change'))\n",
    "        for evaluation in results[label_1][scope]:\n",
    "            val = round(results[label_1][scope][evaluation][metric], 2)\n",
    "            val_cmp = round(results[label_2][scope][evaluation][metric], 2)\n",
    "            \n",
    "            percentage_change = (val_cmp - val) / val * 100 if val != 0 else np.NaN\n",
    "            print(f'{evaluation:>8}: {val:>6.2f} {val_cmp:>6.2f} {percentage_change:>8.2f} %')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d495b4ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "# GENERATION OF (LaTeX) TABLES FROM THE RESULTS\n",
    "\n",
    "# The table displays the evaluation methods in the columns and the methods and target_snr values in the rows.\n",
    "# With the exception of the last experiment, all evaluation metrics (SNR improvement, RMSE, pcorr and CPU time) are displayed. \n",
    "# For the last experiment only the pcorr and RMSE are displayed. \n",
    "\n",
    "# The path to the respective results file\n",
    "filename = 'results/3/improvement_4/results.json'\n",
    "results = json.load(open(filename))\n",
    "\n",
    "# The respective target_snr values from the config of the results file that should appear in the table\n",
    "# target_snrs_print = ['ECG Noise', 'ECG Predicted']\n",
    "# target_snrs_print = ['all']\n",
    "target_snrs_print = ['-1', '0', '2', '5', '7', 'all']\n",
    "\n",
    "last_exp = target_snrs_print == ['ECG Noise', 'ECG Predicted']\n",
    "\n",
    "header = [[method, scope] for method in results.keys() for scope in target_snrs_print]\n",
    "column_names = pd.DataFrame(header, columns=['Method', ''])\n",
    "\n",
    "rows = [[], [], [], []]\n",
    "for method, scope in header:\n",
    "        rows[0].append(f\"${results[method][scope]['snr_imp']['mean']:.2f} \\pm {results[method][scope]['snr_imp']['sd']:.2f}$\")\n",
    "        rows[1].append(f\"${results[method][scope]['rmse']['mean']:.2f} \\pm {results[method][scope]['rmse']['sd']:.2f}$\")\n",
    "        rows[2].append(f\"${results[method][scope]['pcorr']['mean']:.2f} \\pm {results[method][scope]['pcorr']['sd']:.2f}$\")\n",
    "        rows[3].append(f\"${results[method][scope]['cpu_time']['mean']*1000:.2f} \\pm {results[method][scope]['cpu_time']['sd']*1000:.2f}$\")\n",
    "\n",
    "if last_exp:\n",
    "        rows = rows[1:3] # Show only pcorr and RMSE for last experiment\n",
    "\n",
    "columns = pd.MultiIndex.from_frame(column_names)\n",
    "index = ['RMSE', 'pcorr'] if last_exp else ['$SNR_{imp}$', 'RMSE', 'pcorr', 'CPU Time (ms)'] \n",
    "\n",
    "df = pd.DataFrame(rows, columns=columns, index=index)\n",
    "df = df.T # Transpose the dataframe so that the evaluation methods are in the columns and the methods and target_snr values in the rows.\n",
    "display(df)\n",
    "\n",
    "# Print LaTeX code\n",
    "print(df.style.to_latex(clines='skip-last;data', multirow_align='c', sparse_columns=True, sparse_index=True, hrules=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41ef1086",
   "metadata": {},
   "outputs": [],
   "source": [
    "# VISUALISATION FOR WGN NOISE\n",
    "\n",
    "# Prints a figure with only the noise and a figure with the added noise to the signal together with the clean signal itself\n",
    "\n",
    "directory = os.path.join('results', 'noise_viz')\n",
    "if not os.path.exists(directory):\n",
    "    os.makedirs(directory)\n",
    "\n",
    "# Select a random sample from the test dataset\n",
    "i = random.choice(range(0, len(test_dataset)))\n",
    "print(f'Sample index: {i}')\n",
    "\n",
    "_, data = test_dataset[i]\n",
    "\n",
    "data = data.numpy()\n",
    "\n",
    "# Generate WGN noise with the specified target_snr\n",
    "target_snr = 15\n",
    "print(f'Target SNR: {target_snr}')\n",
    "\n",
    "data_avg_watts = np.mean(data ** 2)\n",
    "data_avg_db = 10 * np.log10(data_avg_watts)\n",
    "\n",
    "noise_avg_db = data_avg_db - target_snr\n",
    "noise_avg_watts = 10 ** (noise_avg_db / 10)\n",
    "\n",
    "mean_noise = 0\n",
    "noise = np.random.normal(mean_noise, np.sqrt(noise_avg_watts), len(data))\n",
    "data_noise = data + noise\n",
    "print(f'Acheived SNR: {snr(data, data_noise)}')\n",
    "\n",
    "# The sample range that should be visualised\n",
    "time_from = 0\n",
    "time_to = 512 \n",
    "\n",
    "plt.figure(figsize=(14, 4))\n",
    "\n",
    "x = np.array(range(time_to)[time_from:time_to]) / sampling_rate\n",
    "plt.plot(x, noise[time_from:time_to], color='#087ffb')\n",
    "plt.xticks(x[::25])\n",
    "\n",
    "plt.title('Noise generated using the WGN model')\n",
    "plt.xlabel('Time (s)')\n",
    "plt.ylabel(r'Acceleration')\n",
    "plt.ylim(bottom=-1, top=1)\n",
    "\n",
    "plt.savefig(os.path.join(directory, 'noise_wgn.png'), format='png', bbox_inches='tight')\n",
    "plt.show()\n",
    "\n",
    "\n",
    "plt.figure(figsize=(14, 4))\n",
    "\n",
    "x = np.array(range(time_to)[time_from:time_to]) / sampling_rate\n",
    "plt.plot(x, data[time_from:time_to], color='#087ffb', label='clean signal')\n",
    "plt.plot(x, data_noise[time_from:time_to], color='#ea002e', label='clean signal with WGN')\n",
    "plt.xticks(x[::25])\n",
    "\n",
    "plt.title('SCG signal with added WGN noise')\n",
    "plt.xlabel('Time (s)')\n",
    "plt.ylabel(r'Acceleration')\n",
    "plt.ylim(bottom=-1, top=1)\n",
    "\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47abe456",
   "metadata": {},
   "outputs": [],
   "source": [
    "# VISUALISATION FOR MULTIPLE WGN NOISE SAMPLES\n",
    "\n",
    "# Prints a figure with multiple samples of added noise to the clean signal together with the clean signal itself\n",
    "\n",
    "directory = os.path.join('results', 'noise_viz')\n",
    "if not os.path.exists(directory):\n",
    "    os.makedirs(directory)\n",
    "    \n",
    "color_map = ['#2c9bda', '#0065b0', '#ea002e']\n",
    "\n",
    "# Select a random sample from the test dataset\n",
    "i = random.choice(range(0, len(test_dataset)))\n",
    "i = 2279\n",
    "print(f'Sample index: {i}')\n",
    "\n",
    "_, data = test_dataset[i]\n",
    "\n",
    "data = data.numpy()\n",
    "\n",
    "# Generate WGN noise for the specified target_snrs\n",
    "target_snrs = [7, 15]\n",
    "noise = []\n",
    "\n",
    "mean_noise = 0\n",
    "\n",
    "data_avg_watts = np.mean(data ** 2)\n",
    "data_avg_db = 10 * np.log10(data_avg_watts)\n",
    "\n",
    "# Create the noise for the specified target_snrs\n",
    "for target_snr in target_snrs:\n",
    "    noise_avg_db = data_avg_db - target_snr\n",
    "    noise_avg_watts = 10 ** (noise_avg_db / 10)\n",
    "    noise.append(np.random.normal(mean_noise, np.sqrt(noise_avg_watts), len(data)))\n",
    "\n",
    "# The sample range that should be visualised\n",
    "time_from = 0\n",
    "time_to = 128\n",
    "\n",
    "plt.figure(figsize=(14, 4))\n",
    "x = np.array(range(time_to)[time_from:time_to]) / sampling_rate\n",
    "\n",
    "# Plot the clean signal\n",
    "plt.plot(x, data[time_from:time_to], label='clean', color=color_map[0], linestyle='dashed')\n",
    "\n",
    "# Plot the clean signals with added noise for every noise sample from the `noise` array\n",
    "for i, noise_record in enumerate(noise):\n",
    "    data_noise = data + noise_record\n",
    "    plt.plot(x, data_noise[time_from:time_to], label=f'{target_snrs[i]} dB', color=color_map[i+1])\n",
    "\n",
    "plt.xticks(x[::25])\n",
    "\n",
    "plt.title('Additive Noise Generated Using the WGN model')\n",
    "plt.xlabel('Time (s)')\n",
    "plt.ylabel(r'Acceleration')\n",
    "plt.ylim(bottom=-1, top=1)\n",
    "plt.legend()\n",
    "\n",
    "plt.savefig(os.path.join(directory, 'noise_multi.png'), format='png', bbox_inches='tight')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8efba1d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# VISUALISATIONS FOR MOTION NOISE\n",
    "\n",
    "# Prints one figure with only the generated noise and one with the added noise on a random clean SCG sample together with this clean sample itself\n",
    "\n",
    "directory = os.path.join('results', 'noise_viz')\n",
    "if not os.path.exists(directory):\n",
    "    os.makedirs(directory)\n",
    "\n",
    "# Select a random sample from the test dataset\n",
    "i = random.choice(range(0, len(test_dataset)))\n",
    "i = 6515\n",
    "print(f'Sample index: {i}')\n",
    "\n",
    "_, data = test_dataset[i]\n",
    "\n",
    "data = data.numpy()\n",
    "\n",
    "# Randomly select `len(data)` consecutive samples from the processed noise dataset.\n",
    "num_samples = len(data)\n",
    "rows = range(test_dataset.dataset.data_motion_noise.shape[0])\n",
    "index_start = random.randint(rows.start, rows.stop - num_samples)\n",
    "index_start = 215977\n",
    "print(f'Noise index: {index_start}')\n",
    "\n",
    "noise = test_dataset.dataset.data_motion_noise.iloc[index_start:index_start + num_samples]\n",
    "\n",
    "# Only add 80 % of noise to the data.\n",
    "noise_data = data + 0.8 * noise\n",
    "\n",
    "print(f'SNR: {snr(data, noise_data)}')\n",
    "\n",
    "# The sample range that should be visualised\n",
    "time_from = 0\n",
    "time_to = 512 \n",
    "\n",
    "# Figure with only the noise\n",
    "plt.figure(figsize=(14, 4))\n",
    "\n",
    "x = np.array(range(time_to)[time_from:time_to]) / sampling_rate\n",
    "plt.plot(x, noise.iloc[time_from:time_to], color='#087ffb')\n",
    "plt.xticks(x[::25])\n",
    "\n",
    "plt.title('Acceleration on the z-axis generated by walking')\n",
    "plt.xlabel('Time (s)')\n",
    "plt.ylabel(r'Acceleration')\n",
    "plt.ylim(bottom=-1, top=1)\n",
    "\n",
    "plt.savefig(os.path.join(directory, 'noise_walking.png'), format='png', bbox_inches='tight')\n",
    "plt.show()\n",
    "\n",
    "\n",
    "# Figure with the added noise to the clean signal and the clean signal itself\n",
    "plt.figure(figsize=(14, 4))\n",
    "\n",
    "x = np.array(range(time_to)[time_from:time_to]) / sampling_rate\n",
    "plt.plot(x, data[time_from:time_to], color='#087ffb', label='clean signal')\n",
    "plt.plot(x, noise_data.iloc[time_from:time_to], color='#ea002e', label='clean signal with walking noise')\n",
    "plt.xticks(x[::25])\n",
    "\n",
    "plt.title('SCG Signal with Added Walking Noise')\n",
    "plt.xlabel('Time (s)')\n",
    "plt.ylabel(r'Acceleration')\n",
    "plt.ylim(bottom=-1, top=1)\n",
    "\n",
    "plt.legend()\n",
    "\n",
    "plt.savefig(os.path.join(directory, 'noise_added_walking.png'), format='png', bbox_inches='tight')\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "thesis-project-i_Sg-Li3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.6 (main, Nov 10 2011, 15:00:00) [GCC 12.2.0]"
  },
  "vscode": {
   "interpreter": {
    "hash": "e6717f947e3cb42ac9683eac40a7d814e21ca3f0fa83f7d2ca8c93372cd6a329"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}