diff --git a/brainiak/eventseg/event.py b/brainiak/eventseg/event.py index 4373addd9..ba3f75d58 100644 --- a/brainiak/eventseg/event.py +++ b/brainiak/eventseg/event.py @@ -18,12 +18,18 @@ learning the signature activity pattern of each event, the model can then be applied to other datasets to identify a corresponding sequence of events. -Full details are available in the bioRxiv preprint: +Full details are available in: Christopher Baldassano, Janice Chen, Asieh Zadbood, Jonathan W Pillow, Uri Hasson, Kenneth A Norman Discovering event structure in continuous narrative perception and memory Neuron, Volume 95, Issue 3, 709 - 721.e5 -http://www.cell.com/neuron/abstract/S0896-6273(17)30593-7 +https://doi.org/10.1016/j.neuron.2017.06.041 + +This class also extends the model described in the Neuron paper, by allowing +transition matrices that are composed of multiple separate chains of events +rather than a single linear path. This allows a model to contain patterns for +multiple event sequences (e.g. narratives), and fit probabilities along each of +these chains on a new, unlabeled timeseries. """ # Authors: Chris Baldassano and Cătălin Iordan (Princeton University) @@ -60,6 +66,10 @@ class EventSegment(BaseEstimator): n_iter: int : default 500 Maximum number of steps to run during fitting + event_chains: ndarray with length = n_events + Array with unique value for each separate chain of events, each linked + in the order they appear in the array + Attributes ---------- p_start, p_end: length n_events+1 ndarray @@ -86,10 +96,14 @@ def _default_var_schedule(step): def __init__(self, n_events=2, step_var=_default_var_schedule, - n_iter=500): + n_iter=500, event_chains=None): self.n_events = n_events self.step_var = step_var self.n_iter = n_iter + if event_chains is None: + self.event_chains = np.zeros(n_events) + else: + self.event_chains = event_chains def fit(self, X, y=None): """Learn a segmentation on training data @@ -252,46 +266,54 @@ def _forward_backward(self, logprob): log_beta = np.zeros((t, self.n_events + 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 - p_start = np.zeros((1, self.n_events + 1)) - p_start[0, 0] = 1 - p_trans = (self.n_events-1)/t - P = np.vstack((np.hstack(( - (1 - p_trans) * np.diag(np.ones(self.n_events)) - + p_trans * np.diag(np.ones(self.n_events - 1), 1), - np.append(np.zeros((self.n_events - 1, 1)), [[p_trans]], axis=0))), - np.append(np.zeros((1, self.n_events)), [[1]], - axis=1))) - p_end = np.zeros((1, self.n_events + 1)) - p_end[0, -2] = 1 + self.p_start = np.zeros(self.n_events + 1) + self.p_end = np.zeros(self.n_events + 1) + self.P = np.zeros((self.n_events + 1, self.n_events + 1)) + label_ind = np.unique(self.event_chains, return_inverse=True)[1] + n_chains = np.max(label_ind) + 1 + + # For each chain of events, link them together and then to sink state + for c in range(n_chains): + chain_ind = np.nonzero(label_ind == c)[0] + self.p_start[chain_ind[0]] = 1 / n_chains + self.p_end[chain_ind[-1]] = 1 / n_chains + + p_trans = (len(chain_ind) - 1) / t + if p_trans >= 1: + raise ValueError('Too few timepoints') + for i in range(len(chain_ind)): + self.P[chain_ind[i], chain_ind[i]] = 1 - p_trans + if i < len(chain_ind) - 1: + self.P[chain_ind[i], chain_ind[i+1]] = p_trans + else: + self.P[chain_ind[i], -1] = p_trans + self.P[-1, -1] = 1 # Forward pass for i in range(t): if i == 0: - log_alpha[0, :] = self._log(p_start) + logprob[0, :] + log_alpha[0, :] = self._log(self.p_start) + logprob[0, :] else: log_alpha[i, :] = self._log(np.exp(log_alpha[i - 1, :]) - .dot(P)) + logprob[i, :] + .dot(self.P)) + logprob[i, :] log_scale[i] = np.logaddexp.reduce(log_alpha[i, :]) log_alpha[i] -= log_scale[i] # Backward pass - log_beta[-1, :] = self._log(p_end) - log_scale[-1] + log_beta[-1, :] = self._log(self.p_end) - log_scale[-1] for i in reversed(range(t - 1)): obs_weighted = log_beta[i + 1, :] + logprob[i + 1, :] offset = np.max(obs_weighted) log_beta[i, :] = offset + self._log( - np.exp(obs_weighted - offset).dot(P.T)) - log_scale[i] + np.exp(obs_weighted - offset).dot(self.P.T)) - log_scale[i] # 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(p_end), axis=1) + log_alpha[-1, :] + log_scale[-1] + self._log(self.p_end)) log_gamma = log_gamma[:, :-1] @@ -366,11 +388,14 @@ def find_events(self, testing_data, var=None, scramble=False): if var is None: if not hasattr(self, 'event_var_'): - raise NotFittedError(("The event patterns must first be set " - "by fit() or set_event_patterns()")) + raise NotFittedError(("Event variance must be provided, if " + "not previously set by fit()")) else: var = self.event_var_ + if not hasattr(self, 'event_pat_'): + raise NotFittedError(("The event patterns must first be set " + "by fit() or set_event_patterns()")) if scramble: mean_pat = self.event_pat_[:, np.random.permutation(self.n_events)] else: @@ -440,3 +465,27 @@ def calc_weighted_event_var(self, D, weights, event_pat): np.sum(weights[nz, e])) ev_var = ev_var / D.shape[1] return ev_var + + def model_prior(self, t): + """Returns the prior probability of the HMM + + Runs forward-backward without any data, showing the prior distribution + of the model (for comparison with a posterior). + + Parameters + ---------- + t: int + Number of timepoints + + Returns + ------- + segments : time by event ndarray + segments[t,e] = prior probability that timepoint t is in event e + + test_ll : float + Log-likelihood of model (data-independent term)""" + + lg, test_ll = self._forward_backward(np.zeros((t, self.n_events))) + segments = np.exp(lg) + + return segments, test_ll diff --git a/docs/newsfragments/353.feature b/docs/newsfragments/353.feature new file mode 100644 index 000000000..cb42664a8 --- /dev/null +++ b/docs/newsfragments/353.feature @@ -0,0 +1,2 @@ +Added event_chains option and model_prior function to eventseg + diff --git a/tests/eventseg/test_event.py b/tests/eventseg/test_event.py index dd60134ea..a82de27bd 100644 --- a/tests/eventseg/test_event.py +++ b/tests/eventseg/test_event.py @@ -1,9 +1,12 @@ -import brainiak.eventseg.event +from brainiak.eventseg.event import EventSegment +from scipy.special import comb import numpy as np +import pytest +from sklearn.exceptions import NotFittedError def test_create_event_segmentation(): - es = brainiak.eventseg.event.EventSegment(5) + es = EventSegment(5) assert es, "Invalid EventSegment instance" @@ -11,7 +14,7 @@ def test_fit_shapes(): K = 5 V = 3 T = 10 - es = brainiak.eventseg.event.EventSegment(K, n_iter=2) + es = EventSegment(K, n_iter=2) sample_data = np.random.rand(V, T) es.fit(sample_data.T) @@ -29,21 +32,41 @@ def test_fit_shapes(): assert np.isclose(np.sum(test_segments, axis=1), np.ones(T2)).all(), \ "Segmentation from find_events not correctly normalized" + es_invalid = EventSegment(K) + with pytest.raises(ValueError, message="T < K should cause error"): + es_invalid.model_prior(K-1) + with pytest.raises(ValueError, message="#Events < K should cause error"): + es_invalid.set_event_patterns(np.zeros((V, K-1))) + def test_simple_boundary(): - es = brainiak.eventseg.event.EventSegment(2) - sample_data = np.asarray([[1, 1, 1, 0, 0, 0, 0], [0, 0, 0, 1, 1, 1, 1]]) + es = EventSegment(2) + random_state = np.random.RandomState(0) + + sample_data = np.array([[1, 1, 1, 0, 0, 0, 0], [0, 0, 0, 1, 1, 1, 1]]) + \ + random_state.rand(2, 7) * 10 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" + events_predict = es.predict(sample_data.T) + assert np.array_equal(events_predict, [0, 0, 0, 1, 1, 1, 1]), \ + "Error in predict interface" + def test_event_transfer(): - es = brainiak.eventseg.event.EventSegment(2) - es.set_event_patterns(np.asarray([[1, 0], [0, 1]])) + es = EventSegment(2) sample_data = np.asarray([[1, 1, 1, 0, 0, 0, 0], [0, 0, 0, 1, 1, 1, 1]]) + + with pytest.raises(NotFittedError, message="Should need to set variance"): + seg = es.find_events(sample_data.T)[0] + + with pytest.raises(NotFittedError, message="Should need to set patterns"): + seg = es.find_events(sample_data.T, np.asarray([1, 1]))[0] + + es.set_event_patterns(np.asarray([[1, 0], [0, 1]])) seg = es.find_events(sample_data.T, np.asarray([1, 1]))[0] events = np.argmax(seg, axis=1) @@ -52,7 +75,7 @@ def test_event_transfer(): def test_weighted_var(): - es = brainiak.eventseg.event.EventSegment(2) + es = EventSegment(2) D = np.zeros((8, 4)) for t in range(4): @@ -77,7 +100,7 @@ def test_weighted_var(): def test_sym(): - es = brainiak.eventseg.event.EventSegment(4) + es = EventSegment(4) evpat = np.repeat(np.arange(10).reshape(-1, 1), 4, axis=1) es.set_event_patterns(evpat) @@ -88,3 +111,45 @@ def test_sym(): # Check that events 1-4 and 2-3 are symmetric assert np.all(np.isclose(ev[:, :2], np.fliplr(np.flipud(ev[:, 2:])))),\ "Fit with constant data is not symmetric" + + +def test_chains(): + es = EventSegment(5, event_chains=np.array(['A', 'A', 'B', 'B', 'B'])) + + es.set_event_patterns(np.array([[1, 1, 0, 0, 0], + [0, 0, 1, 1, 1]])) + sample_data = np.array([[0, 0, 0], [1, 1, 1]]) + seg = es.find_events(sample_data.T, 0.1)[0] + + ev = np.nonzero(seg > 0.99)[1] + assert np.array_equal(ev, [2, 3, 4]),\ + "Failed to fit with multiple chains" + + +def test_prior(): + K = 10 + T = 100 + + es = EventSegment(K) + mp = es.model_prior(T)[0] + + p_bound = np.zeros((T, K-1)) + norm = comb(T-1, K-1) + for t in range(T-1): + for k in range(K-1): + # See supplementary material of Neuron paper + # https://doi.org/10.1016/j.neuron.2017.06.041 + p_bound[t+1, k] = comb(t, k) * comb(T-t-2, K-k-2) / norm + p_bound = np.cumsum(p_bound, axis=0) + + mp_gt = np.zeros((T, K)) + for k in range(K): + if k == 0: + mp_gt[:, k] = 1 - p_bound[:, 0] + elif k == K - 1: + mp_gt[:, k] = p_bound[:, k-1] + else: + mp_gt[:, k] = p_bound[:, k-1] - p_bound[:, k] + + assert np.all(np.isclose(mp, mp_gt)),\ + "Prior does not match analytic solution"