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"