From 8f434e0d6f95b90899c4206a7b1266b26825eaa5 Mon Sep 17 00:00:00 2001 From: chris33 Date: Sat, 16 Jul 2016 20:35:40 -0400 Subject: [PATCH] Added Event Segmentation class, along with example and tests Making changes based on pull request comments, including implementing the sklearn Estimator interface Improved simulated data example, added sklearn validation test, added correctness test. Final commit before rebase. Adding nose requirement and fixing docstrings Removed sklearn tests Fixed functions that were modifying their arguments, and added clarifying comments Fixed docstrings --- brainiak/eventseg/__init__.py | 14 ++ brainiak/eventseg/event.py | 369 ++++++++++++++++++++++++++++ examples/eventseg/requirements.txt | 1 + examples/eventseg/simulated_data.py | 101 ++++++++ tests/eventseg/test_event.py | 38 +++ 5 files changed, 523 insertions(+) create mode 100644 brainiak/eventseg/__init__.py create mode 100644 brainiak/eventseg/event.py create mode 100644 examples/eventseg/requirements.txt create mode 100644 examples/eventseg/simulated_data.py create mode 100644 tests/eventseg/test_event.py diff --git a/brainiak/eventseg/__init__.py b/brainiak/eventseg/__init__.py new file mode 100644 index 000000000..5247d8d55 --- /dev/null +++ b/brainiak/eventseg/__init__.py @@ -0,0 +1,14 @@ +# Copyright 2016 Princeton University +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Event segmentation of continuous data + event transfer between datasets.""" diff --git a/brainiak/eventseg/event.py b/brainiak/eventseg/event.py new file mode 100644 index 000000000..7a46f909b --- /dev/null +++ b/brainiak/eventseg/event.py @@ -0,0 +1,369 @@ +# Copyright 2016 Princeton University +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Event segmentation using a Hidden Markov Model + +Given an ROI timeseries, this class uses an annealed fitting procedure to +segment the timeseries into events with stable activity patterns. After +learning the signature activity pattern of each event, the model can then be +applied to other datasets to identify a corresponding sequence of events. +""" + +# Authors: Chris Baldassano and Cătălin Iordan (Princeton University) + +import numpy as np +from scipy import stats +import logging +import copy +from sklearn.base import BaseEstimator +from sklearn.utils.validation import check_is_fitted, check_array + +logger = logging.getLogger(__name__) + +__all__ = [ + "EventSegment", +] + + +class EventSegment(BaseEstimator): + """Class for event segmentation of continuous fMRI data + + Parameters + ---------- + n_events: int + Number of segments to learn + + step_var: Callable[[int], float] : default 4 * (0.98 ** (step - 1)) + The Gaussian variance to use during fitting, as a function of the + number of steps. Should decrease slowly over time. + + n_iter: int : default 500 + Maximum number of steps to run during fitting + + Attributes + ---------- + p_start, p_end: length n_events+1 ndarray + initial and final prior distributions over events + + P: n_events+1 by n_events+1 ndarray + HMM transition matrix + + ll_ : ndarray with length = number of training datasets + Log-likelihood for training datasets over the course of training + + segments_: list of (time by event) ndarrays + Learned (soft) segmentation for training datasets + + event_var_ : float + Gaussian variance at the end of learning + + event_pat_ : voxel by event ndarray + Learned mean patterns for each event + """ + + def _default_var_schedule(step): + return 4 * (0.98 ** (step - 1)) + + def __init__(self, n_events=2, + step_var=_default_var_schedule, + n_iter=500): + self.n_events = n_events + self.step_var = step_var + self.n_iter = n_iter + + def fit(self, X, y=None): + """Learn a segmentation on training data + + Fits event patterns and a segmentation to training data. After + running this function, the learned event patterns can be used to + segment other datasets using find_events + + Parameters + ---------- + X: time by voxel ndarray, or a list of such ndarrays + fMRI data to be segmented. If a list is given, then all datasets + are segmented simultaneously with the same event patterns + + y: not used (added to comply with BaseEstimator definition) + + Returns + ------- + self: the EventSegment object + """ + + X = copy.deepcopy(X) + if type(X) is not list: + X = check_array(X) + X = [X] + + n_train = len(X) + for i in range(n_train): + X[i] = X[i].T + + n_dim = X[0].shape[0] + for i in range(n_train): + assert (X[i].shape[0] == n_dim) + + # Double-check that data is z-scored in time + for i in range(n_train): + X[i] = stats.zscore(X[i], axis=1, ddof=1) + + # Set up transition matrix, with final sink state + # For transition matrix of this form, the transition probability has + # no impact on the final solution, since all valid paths must take + # the same number of transitions + self.p_start = np.zeros((1, self.n_events + 1)) + self.p_start[0, 0] = 1 + self.P = np.vstack((np.hstack(( + 0.5 * np.diag(np.ones(self.n_events)) + + 0.5 * np.diag(np.ones(self.n_events - 1), 1), + np.append(np.zeros((self.n_events - 1, 1)), [[0.5]], axis=0))), + np.append(np.zeros((1, self.n_events)), [[1]], axis=1))) + self.p_end = np.zeros((1, self.n_events + 1)) + self.p_end[0, -2] = 1 + + self.classes_ = np.arange(self.n_events) + + # Initialize variables for fitting + log_gamma = [] + for i in range(n_train): + log_gamma.append(np.zeros((X[i].shape[1], + self.n_events))) + step = 1 + best_ll = float("-inf") + self.ll_ = np.empty((0, n_train)) + while step <= self.n_iter: + iteration_var = self.step_var(step) + + # Based on the current segmentation, compute the mean pattern + # for each event + seg_prob = [np.exp(lg) / np.sum(np.exp(lg), axis=0) + for lg in log_gamma] + mean_pat = np.empty((n_train, n_dim, self.n_events)) + for i in range(n_train): + mean_pat[i, :, :] = X[i].dot(seg_prob[i]) + mean_pat = np.mean(mean_pat, axis=0) + + # Based on the current mean patterns, compute the event + # segmentation + self.ll_ = np.append(self.ll_, np.empty((1, n_train)), axis=0) + for i in range(n_train): + logprob = self._logprob_obs(X[i], + mean_pat, iteration_var) + log_gamma[i], self.ll_[-1, i] = self._forward_backward(logprob) + + # If log-likelihood has started decreasing, undo last step and stop + if np.mean(self.ll_[-1, :]) < best_ll: + self.ll_ = self.ll_[:-1, :] + break + + self.segments_ = [np.exp(lg) for lg in log_gamma] + self.event_var_ = iteration_var + self.event_pat_ = mean_pat + best_ll = np.mean(self.ll_[-1, :]) + logger.debug("Fitting step %d, LL=%f", step, best_ll) + + step += 1 + + return self + + def _logprob_obs(self, data, mean_pat, var): + """Log probability of observing each timepoint under each event model + + Computes the log probability of each observed timepoint being + generated by the Gaussian distribution for each event pattern + + Parameters + ---------- + data: voxel by time ndarray + fMRI data on which to compute log probabilities + + mean_pat: voxel by event ndarray + Centers of the Gaussians for each event + + var: float or 1D array of length equal to the number of events + Variance of the event Gaussians. If scalar, all events are + assumed to have the same variance + + Returns + ------- + logprob : time by event ndarray + Log probability of each timepoint under each event Gaussian + """ + + n_vox = data.shape[0] + t = data.shape[1] + + # z-score both data and mean patterns in space, so that Gaussians + # are measuring Pearson correlations and are insensitive to overall + # activity changes + data_z = stats.zscore(data, axis=0, ddof=1) + mean_pat_z = stats.zscore(mean_pat, axis=0, ddof=1) + + logprob = np.empty((t, self.n_events)) + + if type(var) is not np.ndarray: + var = var * np.ones(self.n_events) + + for k in range(self.n_events): + logprob[:, k] = -0.5 * n_vox * np.log( + 2 * np.pi * var[k]) - 0.5 * np.sum( + (data_z.T - mean_pat_z[:, k]).T ** 2, axis=0) / var[k] + + logprob /= n_vox + return logprob + + def _forward_backward(self, logprob): + """Runs forward-backward algorithm on observation log probs + + Given the log probability of each timepoint being generated by + each event, run the HMM forward-backward algorithm to find the + probability that each timepoint belongs to each event (based on the + transition priors in p_start, p_end, and P) + + See https://en.wikipedia.org/wiki/Forward-backward_algorithm for + mathematical details + + Parameters + ---------- + logprob : time by event ndarray + Log probability of each timepoint under each event Gaussian + + Returns + ------- + log_gamma : time by event ndarray + Log probability of each timepoint belonging to each event + + ll : float + Log-likelihood of fit + """ + logprob = copy.copy(logprob) + t = logprob.shape[0] + logprob = np.hstack((logprob, float("-inf") * np.ones((t, 1)))) + + # Initialize variables + log_scale = np.zeros(t) + log_alpha = np.zeros((t, self.n_events + 1)) + log_beta = np.zeros((t, self.n_events + 1)) + + # Forward pass + for t in range(t): + if t == 0: + log_alpha[0, :] = self._log(self.p_start) + logprob[0, :] + else: + log_alpha[t, :] = self._log(np.exp(log_alpha[t - 1, :]) + .dot(self.P)) + logprob[t, :] + + log_scale[t] = np.logaddexp.reduce(log_alpha[t, :]) + log_alpha[t] -= log_scale[t] + + # Backward pass + log_beta[-1, :] = self._log(self.p_end) - log_scale[-1] + for t in reversed(range(t - 1)): + obs_weighted = log_beta[t + 1, :] + logprob[t + 1, :] + offset = np.max(obs_weighted) + log_beta[t, :] = offset + self._log( + np.exp(obs_weighted - offset).dot(self.P.T)) - log_scale[t] + + # Combine and normalize + log_gamma = log_alpha + log_beta + log_gamma -= np.logaddexp.reduce(log_gamma, axis=1, keepdims=True) + + ll = np.sum(log_scale[:(t - 1)]) + np.logaddexp.reduce( + log_alpha[-1, :] + log_scale[-1] + self._log(self.p_end), axis=1) + + log_gamma = log_gamma[:, :-1] + + return log_gamma, ll + + def _log(self, x): + """Modified version of np.log that manually sets values <=0 to -inf + + Parameters + ---------- + x: ndarray of floats + Input to the log function + + Returns + ------- + log_ma: ndarray of floats + log of x, with x<=0 values replaced with -inf + """ + log_ma = np.ma.log(x).filled(float("-inf")) + return log_ma + + def find_events(self, testing_data, var=None, scramble=False): + """Applies learned event segmentation to new testing dataset + + After fitting an event segmentation using fit(), + this function finds the same sequence of event patterns in a new + testing dataset. + + Parameters + ---------- + testing_data: timepoint by voxel ndarray + fMRI data to segment based on previously-learned event patterns + + var: float or 1D ndarray of length equal to the number of events + default: uses variance that maximized training log-likelihood + Variance of the event Gaussians. If scalar, all events are + assumed to have the same variance. + + scramble: bool : default False + If true, the order of the learned events are shuffled before + fitting, to give a null distribution + + Returns + ------- + segments : time by event ndarray + The resulting soft segmentation. segments[t,e] = probability + that timepoint t is in event e + + test_ll : float + Log-likelihood of model fit + """ + + if var is None: + var = self.event_var_ + + if scramble: + mean_pat = self.event_pat_[:, np.random.permutation(self.n_events)] + else: + mean_pat = self.event_pat_ + + logprob = self._logprob_obs(testing_data.T, mean_pat, var) + lg, test_ll = self._forward_backward(logprob) + segments = np.exp(lg) + + return segments, test_ll + + def predict(self, X): + """Applies learned event segmentation to new testing dataset + + Alternative function for segmenting a new dataset after using + fit() to learn a sequence of events, to comply with the sklearn + Classifier interface + + Parameters + ---------- + X: timepoint by voxel ndarray + fMRI data to segment based on previously-learned event patterns + + Returns + ------- + Event label for each timepoint + """ + check_is_fitted(self, ["event_pat_", "event_var_"]) + X = check_array(X) + segments, test_ll = self.find_events(X) + return np.argmax(segments, axis=1) diff --git a/examples/eventseg/requirements.txt b/examples/eventseg/requirements.txt new file mode 100644 index 000000000..6ccafc3f9 --- /dev/null +++ b/examples/eventseg/requirements.txt @@ -0,0 +1 @@ +matplotlib diff --git a/examples/eventseg/simulated_data.py b/examples/eventseg/simulated_data.py new file mode 100644 index 000000000..deb090bf9 --- /dev/null +++ b/examples/eventseg/simulated_data.py @@ -0,0 +1,101 @@ +"""Example of finding event segmentations on simulated data + +This code generates simulated datasets that have temporally-clustered +structure (with the same series of latent event patterns). An event +segmentation is learned on the first dataset, and then we try to find the same +series of events in other datasets. We measure how well we find the latent +boundaries and the log-likelihood of the fits, and compare to a null model +in which the event order is randomly shuffled. +""" +import brainiak.eventseg.event +import numpy as np +from scipy import stats +import logging +import matplotlib.pyplot as plt + +logging.basicConfig(level=logging.DEBUG) + + +def generate_event_labels(T, K, length_std): + event_labels = np.zeros(T, dtype=int) + start_TR = 0 + for e in range(K - 1): + length = round( + ((T - start_TR) / (K - e)) * (1 + length_std * np.random.randn())) + length = min(max(length, 1), T - start_TR - (K - e)) + event_labels[start_TR:(start_TR + length)] = e + start_TR = start_TR + length + event_labels[start_TR:] = K - 1 + + return event_labels + + +def generate_data(V, T, event_labels, event_means, noise_std): + simul_data = np.empty((V, T)) + for t in range(T): + simul_data[:, t] = stats.multivariate_normal.rvs( + event_means[:, event_labels[t]], cov=noise_std, size=1) + + simul_data = stats.zscore(simul_data, axis=1, ddof=1) + return simul_data + + +# Parameters for creating small simulated datasets +V = 10 +K = 10 +T = 500 +T2 = 300 + +# Generate the first dataset +np.random.seed(1) +event_means = np.random.randn(V, K) +event_labels = generate_event_labels(T, K, 0.1) +simul_data = generate_data(V, T, event_labels, event_means, 1) + +# Find the events in this dataset +simul_seg = brainiak.eventseg.event.EventSegment(K) +simul_seg.fit(simul_data.T) + +# Generate other datasets with the same underlying sequence of event +# patterns, and try to find matching events +test_loops = 10 +bound_match = np.empty((2, test_loops)) +LL = np.empty((2, test_loops)) +for test_i in range(test_loops): + # Generate data + event_labels2 = generate_event_labels(T2, K, 0.5) + simul_data2 = generate_data(V, T2, event_labels2, event_means, 0.1) + + # Find events matching previously-learned events + gamma, LL[0, test_i] = simul_seg.find_events(simul_data2.T) + est_events2 = np.argmax(gamma, axis=1) + bound_match[0, test_i] = 1 - np.sum(abs(np.diff(event_labels2) - + np.diff(est_events2))) / (2 * K) + + # Run again, but with the order of events shuffled so that it no longer + # corresponds to the training data + gamma, LL[1, test_i] = simul_seg.find_events(simul_data2.T, scramble=True) + est_events2 = np.argmax(gamma, axis=1) + bound_match[1, test_i] = 1 - np.sum(abs(np.diff(event_labels2) - + np.diff(est_events2))) / (2 * K) + +# Across the testing datasets, print how well we identify the true event +# boundaries and the log-likehoods in real vs. shuffled data +print("Boundary match: {:.2} (null: {:.2})".format( + np.mean(bound_match[0, :]), np.mean(bound_match[1, :]))) +print("Log-likelihood: {:.3} (null: {:.3})".format( + np.mean(LL[0, :]), np.mean(LL[1, :]))) + +plt.figure() +plt.subplot(2, 1, 1) +plt.imshow(simul_data2, interpolation='nearest', cmap=plt.cm.bone, + aspect='auto') +plt.xlabel('Timepoints') +plt.ylabel('Voxels') +plt.subplot(2, 1, 2) +gamma, LL[0, test_i] = simul_seg.find_events(simul_data2.T) +est_events2 = np.argmax(gamma, axis=1) +plt.plot(est_events2) +plt.xlabel('Timepoints') +plt.ylabel('Event label') +plt.show() diff --git a/tests/eventseg/test_event.py b/tests/eventseg/test_event.py new file mode 100644 index 000000000..3e3b17a95 --- /dev/null +++ b/tests/eventseg/test_event.py @@ -0,0 +1,38 @@ +import brainiak.eventseg.event +import numpy as np + +def test_create_event_segmentation(): + es = brainiak.eventseg.event.EventSegment(5) + assert es, "Invalid EventSegment instance" + + +def test_fit_shapes(): + K = 5 + V = 3 + T = 10 + es = brainiak.eventseg.event.EventSegment(K, n_iter=2) + sample_data = np.random.rand(V, T) + es.fit(sample_data.T) + + assert es.segments_[0].shape == (T,K), "Segmentation from fit " \ + "has incorrect shape" + assert np.isclose(np.sum(es.segments_[0], axis=1), np.ones(T)).all(), \ + "Segmentation from learn_events not correctly normalized" + + T2 = 15 + sample_data2 = np.random.rand(V, T2) + test_segments, test_ll = es.find_events(sample_data2.T) + + assert test_segments.shape == (T2, K), "Segmentation from find_events " \ + "has incorrect shape" + assert np.isclose(np.sum(test_segments, axis=1), np.ones(T2)).all(), \ + "Segmentation from find_events not correctly normalized" + +def test_simple_boundary(): + es = brainiak.eventseg.event.EventSegment(2, n_iter=10) + sample_data = np.asarray([[1, 1, 1, 0, 0, 0, 0], [0, 0, 0, 1, 1, 1, 1]]) + es.fit(sample_data.T) + + events = np.argmax(es.segments_[0], axis=1) + assert np.array_equal(events, [0, 0, 0, 1, 1, 1, 1]),\ + "Failed to correctly segment two events"