HW3: Hidden Markov Models

This assignment considers the problem of classifying sleep states based on noisy heart rate sensor data, like what you might obtain from an Apple Watch. Sleep researchers use polysomnography (PSG) to obtain ground truth sleep states. PSG combines measurements of your brain (EEG), muscles (EMG), and heart (EKG) to determine which of four sleep states you are in:

1. Awake

2. Core Sleep: Non-REM stages 1 and 2. Heart rate and body temperature drop. Sleep spindles, which are thought to be important for memory consolidation, seen in EEG.

3. Deep Sleep: Non-REM stage 3. Hard to wake up, and you feel groggy if you do. Lowest heart rate.

4. Rapid Eye Movement (REM) Sleep: This is when dreaming occurs. Heart rate increases, similar to in an awake state. Not considered restful sleep.

Of course, measuring EEG, EMG, and EKG while you’re sleeping is a big pain. It would be great if we could predict sleep states using less invasive measures, like what you might obtain with an Apple Watch. That’s the goal of this assignment!

This Assignment

Model: You will use a Hidden Markov Model (HMM) to classify sleep stages based on (synthetic) heart rate measurements.

Algorithm: You will use expectation-maximization (EM) and the forward-backward algorithm to estimate the model parameters and infer the latent states.

Data: We will work with a synthetic dataset modeled after a sleep study by Walch et al (2019). We found it too difficult to make accurate predictions on their real data, so we simulated synthetic data with a stronger relationship between sleep states and heart rate. This way, you should be able to predict sleep states with better accuracy than we were able to achieve with the actual data. Though the data is synthetic, it still serves as a good exercise for learning about these models and algorithms.

References

• Walch, O. (2019). Motion and heart rate from a wrist-worn wearable and labeled sleep from polysomnography (version 1.0.0). PhysioNet. https://doi.org/10.13026/hmhs-py35.

Setup

import matplotlib.pyplot as plt
import pandas as pd
import torch

The data consists of 30 time series, each corresponding to one subject’s sleep data over from a single night. The time series have been concatenated together in the data frame below. The fields are,

• id: the index of the time series (0, …, 29)

• state : the “true” sleep state

  • 0 = “awake”

  • 1 = “core sleep”

  • 2 = “deep sleep”

  • 3 = “REM”

• hr: the measured heart rate. Missing data is marked by NaN.

Each row corresponds to a 30 second time bin. The time series are variable in length, ranging from just shy of 4 hours to over 8 hours. Again, these are simulated data, but we generated the data to look like heart rate traces measured by Walch et al. (2019).

# Load the concatenated data
data = pd.read_csv("https://raw.githubusercontent.com/slinderman/stats305b/winter2024/assignments/hw3/hw3_synth.csv")
data

	id	state	hr
0	0	0	97.378630
1	0	0	97.254448
2	0	0	118.884372
3	0	0	99.038121
4	0	0	98.649347
...	...	...	...
26609	29	1	NaN
26610	29	1	NaN
26611	29	1	NaN
26612	29	1	NaN
26613	29	1	NaN

26614 rows × 3 columns

Part 0: Preprocessing

1. Split the data into separate data frames, one for each subject.

2. Plot a couple of the heart rate time series along with the true underlying state.

3. Plot histograms of heart rate for each state.

Make your plots look good, like you would for a figure in a paper.

# Your code here

Part 1: State Estimation

Now we will use an HMM to estimate the sleep states given the heart rate data.

Notation:

Let:

• \(N=30\) denote the number of subjects

• \(K=4\) denote the number of sleep states

• \(T_i\) denote the number of time bins for subject \(i \in \{1, \ldots, N\}\)

• \(z_{i,t} \in \{1, \ldots K\}\) denote the sleep state for subject \(i\) and time \(t\)

• \(x_{i,t} \in \mathbb{R}\) denote the measured heart rate for subject \(i\) at time \(t\).

• \(\boldsymbol{\pi}_0 \in \Delta_{K-1}\) denote the initial distribution

• \(\mathbf{P} \in [0,1]^{K \times K}\) denote the transition matrix

• \(\mu_k \in \mathbb{R}\) denote the conditional mean of the heart rate in state \(k\)

• \(\sigma^2_k \in \mathbb{R}_+\) denote the conditional variance of the heart rate in state \(k\)

• \(\Theta = \{\boldsymbol{\pi}_0, \mathbf{P}, \{\mu_k, \sigma_k^2\}\}\) denote the set of all model parameters

Model:

We will use a Gaussian HMM of the form,

\[\begin{align*} p(\{\{z_{i,t}, x_{i,t}\}_{t=1}^{T_i}\}_{i=1}^N; \Theta) &= \prod_{i=1}^N \left[ \mathrm{Cat}(z_{i,1} \mid \boldsymbol{\pi}_0) \prod_{t=2}^{T_i} \mathrm{Cat}(z_{i, t} \mid \mathbf{P}_{z_{i,t-1}}) \prod_{t=1}^{T_i} \mathrm{N}(x_{i, t} \mid \mu_{z_{i,t}}, \sigma^2_{z_{i,t}}) \right] \end{align*}\]

where \(\mathbf{P}_{k} \in \Delta_{K-1}\) is the \(k\)-th row of the transition matrix.

Assume that the missing heart rate values are “missing at random,” so that the corresponding likelihood terms can simply be dropped from the joint probability above. That is, replace \(\mathrm{N}(x_{i, t} \mid \mu_{z_{i,t}}, \sigma^2_{z_{i,t}})\) with \(1\) for missing values.

Problem 1a: Estimate the Model Parameters

Hold out the last 5 subjects for evaluation. Using the first 25 subjects as your training data, estimate the mean and variance of the heart rate in each sleep state.

1. Estimate \(\mu_k = \mathbb{E}[X_t \mid Z_t=k]\) and \(\sigma_k^2 = \mathrm{Var}[X_t \mid Z_t=k]\). Assume the subjects are iid so you can pool all the data in your estimates.

2. Estimate the transition matrix \(\mathbf{P} \in [0,1]^{K \times K}\).

Assume the initial state distribution is \(\boldsymbol{\pi}_0 = (1, 0, 0, 0)^\top\). That is, assume the subjects always start in the awake state.

# Your code here

Problem 1b: Implement the Forward-Backward Algorithm

Implement the forward-backward algorithm for inferring the posterior probabilities \(\Pr(z_{i,t}=k \mid x_{i,1:T_i}; \Theta)\) for all \(k=1,\ldots,K\) and \(t=1,\ldots,T_i\). Your function should also return the marginal log likelihood, \(\log p(x_{i,1:T_i}; \Theta)\), which is a byproduct of the forward pass (see lecture notes).

Note: Make sure you use a numerically stable implementation!

# Your code here

Problem 1c: Infer the Sleep States

Using your estimated parameters, \(\hat{\Theta}\), from Problem 1a and your implementation from 1b, compute the posterior probabilities of each sleep state and time step on the held-out subjects’ data. Make a nice plot of your results for one subject, including the computed probabilities as well as the true states.

# Your code here

Problem 1d: Evaluate your classifier

Let

\[\begin{align*} \hat{z}_{i,t} = \arg \max_{k \in [K]} \Pr(z_{i,t}=k \mid x_{i,1:T_i}; \hat{\Theta}) \end{align*}\]

denote the state with the highest posterior marginal probability.

Make a confusion matrix to compare your state estimates with the true states of held-out subjects.

Problem 1e: Discussion

Discuss your results from Problems 1c and 1d. How accurate is your classifer? Do the errors make sense?

---

Your answer here

---

Part 2: Parameter Estimation

In Part 1, you fixed the parameters using estimates derived from training data. Often, we don’t have ground truth state labels and we need to simultaneously estimate the states and the parameters. For this part of the assignment, imagine the true state labels are unknown and all you have access to are the heart rate time series.

Problem 2a: Implement the EM Algorithm

Using your implementation of the forward-backward algorithm from above, implement the expectation-maximization (EM) algorithm to simultaneously infer the latent states \(\{\{z_{i,t}\}_{t=1}^{T_i}\}_{i=1}^N\) and estimate the parameters \(\{\mu_k, \sigma_k^2\}_{k=1}^K\).

To keep things simple, you can assume that you have \(\hat{\mathbf{P}}\) from above and that \(\boldsymbol{\pi}_0 = (1,0,0,0)^\top\) is known. (It’s not hard to estimate these parameters too, but it’s a little tedious so we’ll spare you the trouble.)

Your EM function should return:

• the marginal log likelihoods after each iteration of EM

• the final parameter estimates

• the posterior state probabilities under the final parameter estimates

# Your code here

Problem 2b: Run your Code

Run your EM algorithm on the 25 training subjects. Using your estimated parameters, run the forward-backward algorithm to infer the latent states of the held-out 5 subjects.

• Plot the marginal log likelihood as a function of EM iteration.

• Plot a confusion matrix of true and inferred states on the held-out data

Problem 2c: Model Selection

In the unsupervised setting, we often need to estimate the number of discrete states, \(K\), as well. One way to do that is by comparing the marginal log likelihood of held-out data using parameters estimated on the training data, for a range of \(K\).

Sweep over a range of values of \(K\) and for each value run EM on the training data to estimate your model parameters. Then using your forward-backward algorithm to compute the marginal log likelihood of the held-out data. Plot the held-out marginal log likelihood as a function of \(K\).

Note: for this problem you can use a dummy transition matrix that simply imposes a “sticky” prior. Let \(\mathbf{P} = [[p_{i,j}]]\) where \(p_{i,j} = 0.95\) if \(i=j\) and \(p_{i,j} = \frac{0.05}{K-1}\) if \(i \neq j\).

# Your code here

Problem 2d: Discussion

Discuss your results from Part 2. How do your parameter estimates from EM compare to those estimated from the true state labels in Part 1? How many states would you select based on Problem 2c? Does your result make sense?

---

Your answer here

---

Part 3: Model Criticism

As we said in HW1, applied statistics is an iterative process. We don’t just fit one model and call it a day — we revisit our modeling assumptions in light of our findings and look for ways to improve our fit.

Let’s keep \(K=4\) fixed, since we know there are really four official sleep states. Let’s also continue to only consider variations of HMMs. In this part, suggest and then implement at least one improvement to your model, based on your findings above. Quantify your improvements using classification accuracy on held-out subjects.

Note: You may use the true state labels from the training subjects for this part of the assignment.

# Your code here

Submission Instructions

Formatting: check that your code does not exceed 80 characters in line width. You can set Tools → Settings → Editor → Vertical ruler column to 80 to see when you’ve exceeded the limit.

Converting to PDF The simplest way to convert to PDF is to use the “Print to PDF” option in your browser. Just make sure that your code and plots aren’t cut off, as it may not wrap lines.

Alternatively You can download your notebook in .ipynb format and use the following commands to convert it to PDF. Then run the following command to convert to a PDF:

jupyter nbconvert --to pdf <yourlastname>_hw<number>.ipynb

(Note that for the above code to work, you need to rename your file <yourlastname>_hw<number>.ipynb)

Installing nbconvert:

If you’re using Anaconda for package management,

conda install -c anaconda nbconvert

Upload your .pdf file to Gradescope. Please tag your questions correctly! I.e., for each question, all of and only the relevant sections are tagged.

Please post on Ed or come to OH if there are any other problems submitting the HW.