brainiak · mihaic · Mar 27, 2018 · Mar 26, 2018 · Mar 27, 2018 · mihaic
diff --git 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
@@ -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
@@ -1,17 +1,20 @@
-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"
 
 
 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"
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,2 @@
		Added event_chains option and model_prior function to eventseg