From b8b03471d1f39c296cb8a07b76115d65d328832a Mon Sep 17 00:00:00 2001
From: cbaldassano <chrisb33@gmail.com>
Date: Fri, 10 Jan 2020 12:44:46 -0500
Subject: [PATCH 1/8] Full merge/test

---
 brainiak/eventseg/event.py   | 239 +++++++++++++++++++++++++++++++----
 tests/eventseg/test_event.py |  37 ++++--
 2 files changed, 240 insertions(+), 36 deletions(-)

diff --git a/brainiak/eventseg/event.py b/brainiak/eventseg/event.py
index ba3f75d58..59e8c6b34 100644
--- a/brainiak/eventseg/event.py
+++ b/brainiak/eventseg/event.py
@@ -25,11 +25,20 @@
 Neuron, Volume 95, Issue 3, 709 - 721.e5
 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.
+This class also extends the model described in the Neuron paper:
+1) It allows 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.
+To use this option, pass in an event_chain vector labeling which events
+belong to each chain, define event patterns using set_event_patterns(),
+then fit to a new dataset with find_events.
+
+2) To obtain better fits when the underlying event structure contains
+events that vary substantially in length, the merge_split option allows
+the fit() function to re-distribute events during fitting. The number of
+merge/split proposals is controlled by merge_split_proposals, which
+controls how thorough versus fast the fitting process is.
 """
 
 # Authors: Chris Baldassano and Cătălin Iordan (Princeton University)
@@ -41,6 +50,7 @@
 from sklearn.base import BaseEstimator
 from sklearn.utils.validation import check_is_fitted, check_array
 from sklearn.exceptions import NotFittedError
+import itertools
 
 from . import _utils as utils  # type: ignore
 
@@ -63,13 +73,22 @@ class EventSegment(BaseEstimator):
         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
+    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
 
+    merge_split: bool, default: False
+        Determines whether merge/split proposals are used during fitting with
+        fit(). This can improve fitting performance when events are highly
+        uneven in size, but requires additional time
+
+    merge_split_proposals: int, default: 1
+        Number of merges and splits to consider at each step. Computation time
+        scales as O(proposals^2) so this should usually be a small value
+
     Attributes
     ----------
     p_start, p_end: length n_events+1 ndarray
@@ -96,15 +115,55 @@ def _default_var_schedule(step):
 
     def __init__(self, n_events=2,
                  step_var=_default_var_schedule,
-                 n_iter=500, event_chains=None):
+                 n_iter=500, event_chains=None,
+                 merge_split=False, merge_split_proposals=1):
         self.n_events = n_events
         self.step_var = step_var
         self.n_iter = n_iter
+        self.merge_split = merge_split
+        self.merge_split_proposals = merge_split_proposals
         if event_chains is None:
             self.event_chains = np.zeros(n_events)
         else:
             self.event_chains = event_chains
 
+    def _fit_validate(self, X):
+        """Validate input to fit()
+
+        Validate data passed to fit(). Includes a transpose operation to
+        change the row/column order of X and z-scoring in time.
+
+        Parameters
+        ----------
+        X: time by voxel ndarray, or a list of such ndarrays
+            fMRI data to be segmented
+
+        Returns
+        -------
+        X: list of voxel by time ndarrays
+        """
+        if len(np.unique(self.event_chains)) > 1:
+            raise RuntimeError("Cannot fit chains, use set_event_patterns")
+
+        # Copy X into a list and transpose
+        X = copy.deepcopy(X)
+        if type(X) is not list:
+            X = [X]
+        for i in range(len(X)):
+            X[i] = check_array(X[i])
+            X[i] = X[i].T
+
+        # Check that number of voxels is consistent across datasets
+        n_dim = X[0].shape[0]
+        for i in range(len(X)):
+            assert (X[i].shape[0] == n_dim)
+
+        # Double-check that data is z-scored in time
+        for i in range(len(X)):
+            X[i] = stats.zscore(X[i], axis=1, ddof=1)
+
+        return X
+
     def fit(self, X, y=None):
         """Learn a segmentation on training data
 
@@ -125,29 +184,15 @@ def fit(self, X, y=None):
         self: the EventSegment object
         """
 
-        X = copy.deepcopy(X)
-        if type(X) is not list:
-            X = check_array(X)
-            X = [X]
-
+        X = self._fit_validate(X)
         n_train = len(X)
-        for i in range(n_train):
-            X[i] = X[i].T
-
-        self.classes_ = np.arange(self.n_events)
         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)
+        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)))
+            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))
@@ -167,8 +212,7 @@ def fit(self, X, y=None):
             # 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)
+                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
@@ -176,6 +220,11 @@ def fit(self, X, y=None):
                 self.ll_ = self.ll_[:-1, :]
                 break
 
+            if step > 1 and self.merge_split:
+                curr_ll = np.mean(self.ll_[-1, :])
+                self.ll_[-1, :], log_gamma, mean_pat = \
+                    self._merge_split(X, log_gamma, iteration_var, curr_ll)
+
             self.segments_ = [np.exp(lg) for lg in log_gamma]
             self.event_var_ = iteration_var
             self.event_pat_ = mean_pat
@@ -489,3 +538,141 @@ def model_prior(self, t):
         segments = np.exp(lg)
 
         return segments, test_ll
+
+    def _merge_split(self, X, log_gamma, iteration_var, curr_ll):
+        """Attempt to improve log-likelihood with a merge/split
+
+        The simulated annealing used in fit() is susceptible to getting
+        stuck in a local minimum if there are some very short events. This
+        function attempts to find
+        a) pairs of neighboring events that are highly similar, to merge
+        b) events that can be split into two dissimilar events
+        It then tests to see whether simultaneously merging one of the
+        pairs from (a) and splitting one of the events from (b) can improve
+        the log-likelihood. The number of (a)/(b) pairs tested is determined
+        by the merge_split_proposals class attribute.
+
+        Parameters
+        ----------
+        X: list of voxel by time ndarrays
+            fMRI datasets being fit
+
+        log_gamma : list of time by event ndarrays
+            Log probability of each timepoint belonging to each event,
+            for each dataset
+
+        iteration_var : float
+            Current variance in simulated annealing
+
+        curr_ll: float
+            Log-likelihood of current model
+
+        Returns
+        -------
+        return_ll : ndarray with length equal to length of X
+            Log-likelihood after merge/split (same as curr_ll if no
+            merge/split improved curr_ll)
+
+        return_lg : list of time by event ndarrays
+            Log probability of each timepoint belonging to each event,
+            for each dataset (same as log_gamma if no merge/split
+            improved curr_ll)
+
+        return_mp : voxel by event ndarray
+            Mean patterns of events (after possible merge/split)
+        """
+
+        # Compute current probabilities and mean patterns
+        n_train = len(X)
+        n_dim = X[0].shape[0]
+
+        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)
+
+        # Compute correlations between every pair of neighboring
+        # events and their merged pattern. High values of
+        # merge_corr indicate events that are good candidates for
+        # being merged.
+        # Also, compute correlations between every event and each half of
+        # the event, with the midpoint defined as the point where
+        # the cumulative probability (across time) of being in the
+        # event = 0.5. Low values of split_corr indicate events
+        # that are good candidates for being split.
+        merge_corr = np.zeros(self.n_events)
+        merge_pat = np.empty((n_train, n_dim, self.n_events))
+        split_corr = np.zeros(self.n_events)
+        split_pat = np.empty((n_train, n_dim, 2 * self.n_events))
+        for i, sp in enumerate(seg_prob):  # Iterate over datasets
+            merge = np.zeros((sp.shape[0], sp.shape[1]))
+            split = np.zeros((sp.shape[0], 2 * sp.shape[1]))
+            cs = np.cumsum(sp, axis=0)
+            for e in range(sp.shape[1]):
+                mid = np.where(cs[:, e] >= 0.5)[0][0]
+                split[:mid, 2 * e] = sp[:mid, e] / cs[mid - 1, e]
+                split[mid:, 2 * e + 1] = sp[mid:, e] / (1 - cs[mid - 1, e])
+                merge[:, e] = sp[:, e:(e + 2)].mean(1)
+            merge_pat[i, :, :] = X[i].dot(merge)
+            split_pat[i, :, :] = X[i].dot(split)
+        merge_pat = np.mean(merge_pat, axis=0)
+        split_pat = np.mean(split_pat, axis=0)
+        for e in range(self.n_events):
+            split_corr[e] = np.corrcoef(mean_pat[:, e],
+                                        split_pat[:, (2 * e):(2 * e + 2)],
+                                        rowvar=False)[0, 1:3].max()
+            merge_corr[e] = np.corrcoef(merge_pat[:, e],
+                                        mean_pat[:, e:(e + 2)],
+                                        rowvar=False)[0, 1:3].min()
+        merge_corr = merge_corr[:-1]
+
+        # Find best merge/split candidates
+        best_merge = np.flipud(np.argsort(merge_corr))
+        best_merge = best_merge[:self.merge_split_proposals]
+        best_split = np.argsort(split_corr)
+        best_split = best_split[:self.merge_split_proposals]
+
+        # For every pair of merge/split candidates, attempt the merge/split
+        # and measure the log-likelihood. If any are better than curr_ll,
+        # accept this best merge/split
+        mean_pat_last = mean_pat.copy()
+        return_ll = curr_ll
+        return_lg = copy.deepcopy(log_gamma)
+        return_mp = mean_pat.copy()
+        for m_e, s_e in itertools.product(best_merge, best_split):
+            if m_e == s_e or m_e+1 == s_e:
+                # Don't attempt to merge/split same event
+                continue
+
+            # Construct new set of patterns with merge/split
+            mean_pat_ms = np.delete(mean_pat_last,
+                                    [s_e, m_e, m_e + 1], axis=1)
+            s_insert = s_e - 2*(s_e > m_e)
+            m_insert = m_e - (s_e < m_e)
+
+            # Insert in order
+            insert_inds = np.array([s_insert, s_insert, m_insert])
+            insert_pat = np.hstack((split_pat[:, (2 * s_e):(2 * s_e + 2)],
+                                    merge_pat[:, m_e:(m_e+1)]))
+            mean_pat_ms = np.insert(mean_pat_ms,
+                                    insert_inds, insert_pat, axis=1)
+
+            # Measure log-likelihood with these new patterns
+            ll_ms = np.zeros(n_train)
+            log_gamma_ms = list()
+            for i in range(n_train):
+                logprob = self._logprob_obs(X[i],
+                                            mean_pat_ms, iteration_var)
+                lg, ll_ms[i] = self._forward_backward(logprob)
+                log_gamma_ms.append(lg)
+
+            # If better than best ll so far, save to return to fit()
+            if ll_ms.mean() > return_ll:
+                return_mp = mean_pat_ms.copy()
+                return_ll = ll_ms
+                for i in range(n_train):
+                    return_lg[i] = log_gamma_ms[i].copy()
+
+        return return_ll, return_lg, return_mp
diff --git a/tests/eventseg/test_event.py b/tests/eventseg/test_event.py
index 56106014b..a36d8f138 100644
--- a/tests/eventseg/test_event.py
+++ b/tests/eventseg/test_event.py
@@ -33,13 +33,10 @@ def test_fit_shapes():
         "Segmentation from find_events not correctly normalized"
 
     es_invalid = EventSegment(K)
-    with pytest.raises(ValueError):
+    with pytest.raises(ValueError, message="T < K should cause error"):
         es_invalid.model_prior(K-1)
-        # ``with`` block is about to end with no error.
-        pytest.fail("T < K should cause error")
-    with pytest.raises(ValueError):
+    with pytest.raises(ValueError, message="#Events < K should cause error"):
         es_invalid.set_event_patterns(np.zeros((V, K-1)))
-        pytest.fail("#Events < K should cause error")
 
 
 def test_simple_boundary():
@@ -63,13 +60,11 @@ def test_event_transfer():
     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):
+    with pytest.raises(NotFittedError, message="Should need to set variance"):
         seg = es.find_events(sample_data.T)[0]
-        pytest.fail("Should need to set variance")
 
-    with pytest.raises(NotFittedError):
+    with pytest.raises(NotFittedError, message="Should need to set patterns"):
         seg = es.find_events(sample_data.T, np.asarray([1, 1]))[0]
-        pytest.fail("Should need to set patterns")
 
     es.set_event_patterns(np.asarray([[1, 0], [0, 1]]))
     seg = es.find_events(sample_data.T, np.asarray([1, 1]))[0]
@@ -120,10 +115,14 @@ def test_sym():
 
 def test_chains():
     es = EventSegment(5, event_chains=np.array(['A', 'A', 'B', 'B', 'B']))
+    sample_data = np.array([[0, 0, 0], [1, 1, 1]])
+
+    with pytest.raises(RuntimeError,
+                       message="Shouldn't allow fit() with event chains"):
+        seg = es.fit(sample_data.T)[0]
 
     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]
@@ -158,3 +157,21 @@ def test_prior():
 
     assert np.all(np.isclose(mp, mp_gt)),\
         "Prior does not match analytic solution"
+
+
+def test_split_merge():
+    ev = np.array(
+        [0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,
+         3, 3, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4])
+    random_state = np.random.RandomState(0)
+    ev_pat = random_state.rand(5, 10)
+    D = np.zeros((len(ev), 10))
+    for t in range(len(ev)):
+        D[t, :] = ev_pat[ev[t], :] + 0.1*random_state.rand(10)
+
+    hmm_sm = EventSegment(5, merge_split=True, merge_split_proposals=2)
+    hmm_sm.fit(D)
+    hmm_events = np.argmax(hmm_sm.segments_[0], axis=1)
+
+    assert np.all(ev == hmm_events),\
+        "Merge/split fails to find highly uneven events"

From cbea18ce743a9238a7aaf5d3bfab32fd45b60f25 Mon Sep 17 00:00:00 2001
From: cbaldassano <chrisb33@gmail.com>
Date: Fri, 10 Jan 2020 12:49:06 -0500
Subject: [PATCH 2/8] Fixing tests

---
 tests/eventseg/test_event.py | 11 +++++------
 1 file changed, 5 insertions(+), 6 deletions(-)

diff --git a/tests/eventseg/test_event.py b/tests/eventseg/test_event.py
index a36d8f138..b602e3790 100644
--- a/tests/eventseg/test_event.py
+++ b/tests/eventseg/test_event.py
@@ -33,9 +33,9 @@ def test_fit_shapes():
         "Segmentation from find_events not correctly normalized"
 
     es_invalid = EventSegment(K)
-    with pytest.raises(ValueError, message="T < K should cause error"):
+    with pytest.raises(ValueError):
         es_invalid.model_prior(K-1)
-    with pytest.raises(ValueError, message="#Events < K should cause error"):
+    with pytest.raises(ValueError):
         es_invalid.set_event_patterns(np.zeros((V, K-1)))
 
 
@@ -60,10 +60,10 @@ def test_event_transfer():
     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"):
+    with pytest.raises(NotFittedError):
         seg = es.find_events(sample_data.T)[0]
 
-    with pytest.raises(NotFittedError, message="Should need to set patterns"):
+    with pytest.raises(NotFittedError):
         seg = es.find_events(sample_data.T, np.asarray([1, 1]))[0]
 
     es.set_event_patterns(np.asarray([[1, 0], [0, 1]]))
@@ -117,8 +117,7 @@ def test_chains():
     es = EventSegment(5, event_chains=np.array(['A', 'A', 'B', 'B', 'B']))
     sample_data = np.array([[0, 0, 0], [1, 1, 1]])
 
-    with pytest.raises(RuntimeError,
-                       message="Shouldn't allow fit() with event chains"):
+    with pytest.raises(RuntimeError):
         seg = es.fit(sample_data.T)[0]
 
     es.set_event_patterns(np.array([[1, 1, 0, 0, 0],

From 482c6d14c132880c1a879cee4fba79ab507d4605 Mon Sep 17 00:00:00 2001
From: cbaldassano <chrisb33@gmail.com>
Date: Fri, 10 Jan 2020 12:51:56 -0500
Subject: [PATCH 3/8] Fixing tests

---
 tests/eventseg/test_event.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/tests/eventseg/test_event.py b/tests/eventseg/test_event.py
index b602e3790..79825f09f 100644
--- a/tests/eventseg/test_event.py
+++ b/tests/eventseg/test_event.py
@@ -35,8 +35,11 @@ def test_fit_shapes():
     es_invalid = EventSegment(K)
     with pytest.raises(ValueError):
         es_invalid.model_prior(K-1)
+        # ``with`` block is about to end with no error.
+        pytest.fail("T < K should cause error")
     with pytest.raises(ValueError):
         es_invalid.set_event_patterns(np.zeros((V, K-1)))
+        pytest.fail("#Events < K should cause error")
 
 
 def test_simple_boundary():
@@ -62,9 +65,11 @@ def test_event_transfer():
 
     with pytest.raises(NotFittedError):
         seg = es.find_events(sample_data.T)[0]
+        pytest.fail("Should need to set variance")
 
     with pytest.raises(NotFittedError):
         seg = es.find_events(sample_data.T, np.asarray([1, 1]))[0]
+        pytest.fail("Should need to set patterns")
 
     es.set_event_patterns(np.asarray([[1, 0], [0, 1]]))
     seg = es.find_events(sample_data.T, np.asarray([1, 1]))[0]
@@ -119,6 +124,7 @@ def test_chains():
 
     with pytest.raises(RuntimeError):
         seg = es.fit(sample_data.T)[0]
+        pytest.fail("Can't use fit() with event chains")
 
     es.set_event_patterns(np.array([[1, 1, 0, 0, 0],
                                     [0, 0, 1, 1, 1]]))

From 57b643fc9c985fac6b4ab80397a2542e9ce5f112 Mon Sep 17 00:00:00 2001
From: cbaldassano <chrisb33@gmail.com>
Date: Mon, 13 Jan 2020 15:54:13 -0500
Subject: [PATCH 4/8] Simplifying insert code and adding additional test to
 verify that log-likelihood is symmetric

---
 brainiak/eventseg/event.py   | 32 ++++++++++++++++----------------
 tests/eventseg/test_event.py | 21 +++++++++++++++++++++
 2 files changed, 37 insertions(+), 16 deletions(-)

diff --git a/brainiak/eventseg/event.py b/brainiak/eventseg/event.py
index 59e8c6b34..cbc13c4e0 100644
--- a/brainiak/eventseg/event.py
+++ b/brainiak/eventseg/event.py
@@ -215,16 +215,16 @@ def fit(self, X, y=None):
                 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
-
             if step > 1 and self.merge_split:
                 curr_ll = np.mean(self.ll_[-1, :])
                 self.ll_[-1, :], log_gamma, mean_pat = \
                     self._merge_split(X, log_gamma, iteration_var, curr_ll)
 
+            # 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
@@ -647,17 +647,15 @@ def _merge_split(self, X, log_gamma, iteration_var, curr_ll):
                 continue
 
             # Construct new set of patterns with merge/split
-            mean_pat_ms = np.delete(mean_pat_last,
-                                    [s_e, m_e, m_e + 1], axis=1)
-            s_insert = s_e - 2*(s_e > m_e)
-            m_insert = m_e - (s_e < m_e)
-
-            # Insert in order
-            insert_inds = np.array([s_insert, s_insert, m_insert])
-            insert_pat = np.hstack((split_pat[:, (2 * s_e):(2 * s_e + 2)],
-                                    merge_pat[:, m_e:(m_e+1)]))
-            mean_pat_ms = np.insert(mean_pat_ms,
-                                    insert_inds, insert_pat, axis=1)
+            mean_pat_ms = np.delete(mean_pat_last, s_e, axis=1)
+            mean_pat_ms = np.insert(mean_pat_ms, [s_e, s_e],
+                                    split_pat[:, (2 * s_e):(2 * s_e + 2)],
+                                    axis=1)
+            mean_pat_ms = np.delete(mean_pat_ms,
+                                    [m_e + (s_e < m_e), m_e + (s_e < m_e) + 1],
+                                    axis=1)
+            mean_pat_ms = np.insert(mean_pat_ms, m_e + (s_e < m_e),
+                                    merge_pat[:, m_e], axis=1)
 
             # Measure log-likelihood with these new patterns
             ll_ms = np.zeros(n_train)
@@ -674,5 +672,7 @@ def _merge_split(self, X, log_gamma, iteration_var, curr_ll):
                 return_ll = ll_ms
                 for i in range(n_train):
                     return_lg[i] = log_gamma_ms[i].copy()
+                logger.debug("Identified merge %d,%d and split %d",
+                             m_e, m_e+1, s_e)
 
         return return_ll, return_lg, return_mp
diff --git a/tests/eventseg/test_event.py b/tests/eventseg/test_event.py
index 79825f09f..bbe93f53c 100644
--- a/tests/eventseg/test_event.py
+++ b/tests/eventseg/test_event.py
@@ -180,3 +180,24 @@ def test_split_merge():
 
     assert np.all(ev == hmm_events),\
         "Merge/split fails to find highly uneven events"
+
+
+def test_sym_ll():
+    ev = np.array([0, 0, 0, 1, 1, 1, 1, 1, 1, 2, 2])
+    random_state = np.random.RandomState(0)
+    ev_pat = random_state.rand(3, 10)
+    D_forward = np.zeros((len(ev), 10))
+    for t in range(len(ev)):
+        D_forward[t, :] = ev_pat[ev[t], :] + 0.1 * random_state.rand(10)
+    D_backward = np.flip(D_forward, axis=0)
+
+    hmm_forward = EventSegment(3)
+    hmm_forward.set_event_patterns(ev_pat.T)
+    _, ll_forward = hmm_forward.find_events(D_forward, var=1)
+
+    hmm_backward = EventSegment(3)
+    hmm_backward.set_event_patterns(np.flip(ev_pat.T, axis=1))
+    _, ll_backward = hmm_backward.find_events(D_backward, var=1)
+
+    assert (ll_forward == ll_backward),\
+        "Log-likelihood not symmetric forward/backward"

From 4bbda3f24b4a2a99e8c1d0600bf422bf5dc6eb2f Mon Sep 17 00:00:00 2001
From: cbaldassano <chrisb33@gmail.com>
Date: Mon, 13 Jan 2020 16:16:30 -0500
Subject: [PATCH 5/8] Rename merge_split to split_merge for consistency with
 other publications

---
 brainiak/eventseg/event.py   | 26 +++++++++++++-------------
 tests/eventseg/test_event.py |  2 +-
 2 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/brainiak/eventseg/event.py b/brainiak/eventseg/event.py
index cbc13c4e0..b495076fa 100644
--- a/brainiak/eventseg/event.py
+++ b/brainiak/eventseg/event.py
@@ -35,9 +35,9 @@
 then fit to a new dataset with find_events.
 
 2) To obtain better fits when the underlying event structure contains
-events that vary substantially in length, the merge_split option allows
+events that vary substantially in length, the split_merge option allows
 the fit() function to re-distribute events during fitting. The number of
-merge/split proposals is controlled by merge_split_proposals, which
+merge/split proposals is controlled by split_merge_proposals, which
 controls how thorough versus fast the fitting process is.
 """
 
@@ -80,12 +80,12 @@ class EventSegment(BaseEstimator):
         Array with unique value for each separate chain of events, each linked
         in the order they appear in the array
 
-    merge_split: bool, default: False
+    split_merge: bool, default: False
         Determines whether merge/split proposals are used during fitting with
         fit(). This can improve fitting performance when events are highly
         uneven in size, but requires additional time
 
-    merge_split_proposals: int, default: 1
+    split_merge_proposals: int, default: 1
         Number of merges and splits to consider at each step. Computation time
         scales as O(proposals^2) so this should usually be a small value
 
@@ -116,12 +116,12 @@ def _default_var_schedule(step):
     def __init__(self, n_events=2,
                  step_var=_default_var_schedule,
                  n_iter=500, event_chains=None,
-                 merge_split=False, merge_split_proposals=1):
+                 split_merge=False, split_merge_proposals=1):
         self.n_events = n_events
         self.step_var = step_var
         self.n_iter = n_iter
-        self.merge_split = merge_split
-        self.merge_split_proposals = merge_split_proposals
+        self.split_merge = split_merge
+        self.split_merge_proposals = split_merge_proposals
         if event_chains is None:
             self.event_chains = np.zeros(n_events)
         else:
@@ -215,10 +215,10 @@ def fit(self, X, y=None):
                 logprob = self._logprob_obs(X[i], mean_pat, iteration_var)
                 log_gamma[i], self.ll_[-1, i] = self._forward_backward(logprob)
 
-            if step > 1 and self.merge_split:
+            if step > 1 and self.split_merge:
                 curr_ll = np.mean(self.ll_[-1, :])
                 self.ll_[-1, :], log_gamma, mean_pat = \
-                    self._merge_split(X, log_gamma, iteration_var, curr_ll)
+                    self._split_merge(X, log_gamma, iteration_var, curr_ll)
 
             # If log-likelihood has started decreasing, undo last step and stop
             if np.mean(self.ll_[-1, :]) < best_ll:
@@ -539,7 +539,7 @@ def model_prior(self, t):
 
         return segments, test_ll
 
-    def _merge_split(self, X, log_gamma, iteration_var, curr_ll):
+    def _split_merge(self, X, log_gamma, iteration_var, curr_ll):
         """Attempt to improve log-likelihood with a merge/split
 
         The simulated annealing used in fit() is susceptible to getting
@@ -550,7 +550,7 @@ def _merge_split(self, X, log_gamma, iteration_var, curr_ll):
         It then tests to see whether simultaneously merging one of the
         pairs from (a) and splitting one of the events from (b) can improve
         the log-likelihood. The number of (a)/(b) pairs tested is determined
-        by the merge_split_proposals class attribute.
+        by the split_merge_proposals class attribute.
 
         Parameters
         ----------
@@ -630,9 +630,9 @@ def _merge_split(self, X, log_gamma, iteration_var, curr_ll):
 
         # Find best merge/split candidates
         best_merge = np.flipud(np.argsort(merge_corr))
-        best_merge = best_merge[:self.merge_split_proposals]
+        best_merge = best_merge[:self.split_merge_proposals]
         best_split = np.argsort(split_corr)
-        best_split = best_split[:self.merge_split_proposals]
+        best_split = best_split[:self.split_merge_proposals]
 
         # For every pair of merge/split candidates, attempt the merge/split
         # and measure the log-likelihood. If any are better than curr_ll,
diff --git a/tests/eventseg/test_event.py b/tests/eventseg/test_event.py
index bbe93f53c..30b32add5 100644
--- a/tests/eventseg/test_event.py
+++ b/tests/eventseg/test_event.py
@@ -174,7 +174,7 @@ def test_split_merge():
     for t in range(len(ev)):
         D[t, :] = ev_pat[ev[t], :] + 0.1*random_state.rand(10)
 
-    hmm_sm = EventSegment(5, merge_split=True, merge_split_proposals=2)
+    hmm_sm = EventSegment(5, split_merge=True, split_merge_proposals=2)
     hmm_sm.fit(D)
     hmm_events = np.argmax(hmm_sm.segments_[0], axis=1)
 

From e4892a579e8ab6c0ce357c08612dce532635fc0e Mon Sep 17 00:00:00 2001
From: cbaldassano <chrisb33@gmail.com>
Date: Tue, 14 Jan 2020 10:21:37 -0500
Subject: [PATCH 6/8] Tweak to prevent out-of-bounds indexing edge case

---
 brainiak/eventseg/event.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/brainiak/eventseg/event.py b/brainiak/eventseg/event.py
index b495076fa..23bfa3431 100644
--- a/brainiak/eventseg/event.py
+++ b/brainiak/eventseg/event.py
@@ -1,4 +1,4 @@
-#  Copyright 2016 Princeton University
+#  Copyright 2020 Princeton University
 #
 #  Licensed under the Apache License, Version 2.0 (the "License");
 #  you may not use this file except in compliance with the License.
@@ -612,8 +612,10 @@ def _split_merge(self, X, log_gamma, iteration_var, curr_ll):
             cs = np.cumsum(sp, axis=0)
             for e in range(sp.shape[1]):
                 mid = np.where(cs[:, e] >= 0.5)[0][0]
-                split[:mid, 2 * e] = sp[:mid, e] / cs[mid - 1, e]
-                split[mid:, 2 * e + 1] = sp[mid:, e] / (1 - cs[mid - 1, e])
+                cs_first = cs[mid, e] - sp[mid, e]
+                cs_second = 1 - cs_first
+                split[:mid, 2 * e] = sp[:mid, e] / cs_first
+                split[mid:, 2 * e + 1] = sp[mid:, e] / cs_second
                 merge[:, e] = sp[:, e:(e + 2)].mean(1)
             merge_pat[i, :, :] = X[i].dot(merge)
             split_pat[i, :, :] = X[i].dot(split)

From dc0eab8e4e1c3edfc25198b13dced107b6b86116 Mon Sep 17 00:00:00 2001
From: cbaldassano <chrisb33@gmail.com>
Date: Wed, 5 Feb 2020 09:04:42 -0500
Subject: [PATCH 7/8] Adding comments and changing variable names

---
 brainiak/eventseg/event.py | 48 ++++++++++++++++++++++++--------------
 1 file changed, 30 insertions(+), 18 deletions(-)

diff --git a/brainiak/eventseg/event.py b/brainiak/eventseg/event.py
index 23bfa3431..20904f290 100644
--- a/brainiak/eventseg/event.py
+++ b/brainiak/eventseg/event.py
@@ -593,34 +593,40 @@ def _split_merge(self, X, log_gamma, iteration_var, curr_ll):
             mean_pat[i, :, :] = X[i].dot(seg_prob[i])
         mean_pat = np.mean(mean_pat, axis=0)
 
-        # Compute correlations between every pair of neighboring
-        # events and their merged pattern. High values of
-        # merge_corr indicate events that are good candidates for
-        # being merged.
-        # Also, compute correlations between every event and each half of
-        # the event, with the midpoint defined as the point where
-        # the cumulative probability (across time) of being in the
-        # event = 0.5. Low values of split_corr indicate events
-        # that are good candidates for being split.
-        merge_corr = np.zeros(self.n_events)
+        # For each event, merge its probability distribution
+        # with the next event, and also split its probability
+        # distribution at its median into two separate events.
+        # Use these new event probability distributions to compute
+        # merged and split event patterns.
         merge_pat = np.empty((n_train, n_dim, self.n_events))
-        split_corr = np.zeros(self.n_events)
         split_pat = np.empty((n_train, n_dim, 2 * self.n_events))
         for i, sp in enumerate(seg_prob):  # Iterate over datasets
-            merge = np.zeros((sp.shape[0], sp.shape[1]))
-            split = np.zeros((sp.shape[0], 2 * sp.shape[1]))
+            m_evprob = np.zeros((sp.shape[0], sp.shape[1]))
+            s_evprob = np.zeros((sp.shape[0], 2 * sp.shape[1]))
             cs = np.cumsum(sp, axis=0)
             for e in range(sp.shape[1]):
+                # Split distribution at midpoint and normalize each half
                 mid = np.where(cs[:, e] >= 0.5)[0][0]
                 cs_first = cs[mid, e] - sp[mid, e]
                 cs_second = 1 - cs_first
-                split[:mid, 2 * e] = sp[:mid, e] / cs_first
-                split[mid:, 2 * e + 1] = sp[mid:, e] / cs_second
-                merge[:, e] = sp[:, e:(e + 2)].mean(1)
-            merge_pat[i, :, :] = X[i].dot(merge)
-            split_pat[i, :, :] = X[i].dot(split)
+                s_evprob[:mid, 2 * e] = sp[:mid, e] / cs_first
+                s_evprob[mid:, 2 * e + 1] = sp[mid:, e] / cs_second
+
+                # Merge distribution with next event distribution
+                m_evprob[:, e] = sp[:, e:(e + 2)].mean(1)
+
+            # Weight data by distribution to get event patterns
+            merge_pat[i, :, :] = X[i].dot(m_evprob)
+            split_pat[i, :, :] = X[i].dot(s_evprob)
+
+        # Average across datasets
         merge_pat = np.mean(merge_pat, axis=0)
         split_pat = np.mean(split_pat, axis=0)
+
+        # Correlate the current event patterns with the split and
+        # merged patterns
+        merge_corr = np.zeros(self.n_events)
+        split_corr = np.zeros(self.n_events)
         for e in range(self.n_events):
             split_corr[e] = np.corrcoef(mean_pat[:, e],
                                         split_pat[:, (2 * e):(2 * e + 2)],
@@ -631,6 +637,12 @@ def _split_merge(self, X, log_gamma, iteration_var, curr_ll):
         merge_corr = merge_corr[:-1]
 
         # Find best merge/split candidates
+        # A high value of merge_corr indicates that a pair of events are
+        # very similar to their merged pattern, and are good candidates for
+        # being merged.
+        # A low value of split_corr indicates that an event's pattern is
+        # very dissimilar from the patterns in its first and second half,
+        # and is a good candidate for being split.
         best_merge = np.flipud(np.argsort(merge_corr))
         best_merge = best_merge[:self.split_merge_proposals]
         best_split = np.argsort(split_corr)

From 53d05aecf753a2ae60f7ef5098137fcc8f048a0a Mon Sep 17 00:00:00 2001
From: cbaldassano <chrisb33@gmail.com>
Date: Thu, 6 Feb 2020 13:43:51 -0500
Subject: [PATCH 8/8] Adding news fragment

---
 docs/newsfragments/451.feature | 1 +
 1 file changed, 1 insertion(+)
 create mode 100644 docs/newsfragments/451.feature

diff --git a/docs/newsfragments/451.feature b/docs/newsfragments/451.feature
new file mode 100644
index 000000000..bc1a5c1e6
--- /dev/null
+++ b/docs/newsfragments/451.feature
@@ -0,0 +1 @@
+Add eventseg split-merge feature to improve fits when event lengths are highly uneven
\ No newline at end of file