HMM Utils API Reference

class PoissonHMM(PHMM):
    """Nelpy extension of PoissonHMM: Hidden Markov Model with
    independent Poisson emissions.

    Parameters
    ----------
    n_components : int
        Number of states.

    startprob_prior : array, shape (n_components, )
        Initial state occupation prior distribution.

    transmat_prior : array, shape (n_components, n_components)
        Matrix of prior transition probabilities between states.

    algorithm : string, one of the :data:`base.DECODER_ALGORITHMS`
        Decoder algorithm.

    random_state: RandomState or an int seed (0 by default)
        A random number generator instance.

    n_iter : int, optional
        Maximum number of iterations to perform.

    tol : float, optional
        Convergence threshold. EM will stop if the gain in log-likelihood
        is below this value.

    verbose : bool, optional
        When ``True`` per-iteration convergence reports are printed
        to :data:`sys.stderr`. You can diagnose convergence via the
        :attr:`monitor_` attribute.

    params : string, optional
        Controls which parameters are updated in the training
        process.  Can contain any combination of 's' for startprob,
        't' for transmat, 'm' for means and 'c' for covars. Defaults
        to all parameters.

    init_params : string, optional
        Controls which parameters are initialized prior to
        training.  Can contain any combination of 's' for
        startprob, 't' for transmat, 'm' for means and 'c' for covars.
        Defaults to all parameters.

    Attributes
    ----------
    n_features : int
        Dimensionality of the (independent) Poisson emissions.

    monitor_ : ConvergenceMonitor
        Monitor object used to check the convergence of EM.

    transmat_ : array, shape (n_components, n_components)
        Matrix of transition probabilities between states.

    startprob_ : array, shape (n_components, )
        Initial state occupation distribution.

    means_ : array, shape (n_components, n_features)
        Mean parameters for each state.

    extern_ : array, shape (n_components, n_extern)
        Augmented mapping from state space to external variables.

    Examples
    --------
    >>> from nelpy.hmmutils import PoissonHMM
    >>> PoissonHMM(n_components=2)...

    """

    __attributes__ = ["_fs", "_ds", "_unit_ids", "_unit_labels", "_unit_tags"]

    def __init__(
        self,
        *,
        n_components,
        n_iter=None,
        init_params=None,
        params=None,
        random_state=None,
        verbose=False,
    ):
        # assign default parameter values
        if n_iter is None:
            n_iter = 50
        if init_params is None:
            init_params = "stm"
        if params is None:
            params = "stm"

        # TODO: I don't understand why super().__init__ does not work?
        PHMM.__init__(
            self,
            n_components=n_components,
            n_iter=n_iter,
            init_params=init_params,
            params=params,
            random_state=random_state,
            verbose=verbose,
        )

        # initialize BinnedSpikeTrain attributes
        for attrib in self.__attributes__:
            exec("self." + attrib + " = None")

        self._extern_ = None
        self._ds = None
        # self._extern_map = None

        # create shortcuts to super() methods that are overridden in
        # this class
        self._fit = PHMM.fit
        self._score = PHMM.score
        self._score_samples = PHMM.score_samples
        self._predict = PHMM.predict
        self._predict_proba = PHMM.predict_proba
        self._decode = PHMM.decode

        self._sample = PHMM.sample

    def __repr__(self):
        try:
            rep = super().__repr__()
        except Exception:
            warn(
                "couldn't access super().__repr__;"
                " upgrade dependencies to resolve this issue."
            )
            rep = "PoissonHMM"
        if self._extern_ is not None:
            fit_ext = "True"
        else:
            fit_ext = "False"
        try:
            fit = "False"
            if self.means_ is not None:
                fit = "True"
        except AttributeError:
            fit = "False"
        fitstr = "; fit=" + fit + ", fit_ext=" + fit_ext
        return "nelpy." + rep + fitstr

    @property
    def extern_(self):
        """
        Mapping from states to external variables (e.g., position).

        Returns
        -------
        np.ndarray or None
            Array of shape (n_components, n_extern) containing the mapping
            from states to external variables. Returns None if no mapping
            has been learned yet.

        Examples
        --------
        >>> hmm.fit_ext(bst, position_data)
        >>> extern_map = hmm.extern_
        >>> print(f"State 0 maps to position bin {np.argmax(extern_map[0])}")
        """
        if self._extern_ is not None:
            return self._extern_
        else:
            warn("no state <--> external mapping has been learnt yet!")
            return None

    def _get_order_from_transmat(self, start_state=None):
        """Determine a state ordering based on the transition matrix.

        This is a greedy approach, starting at the a priori most probable
        state, and moving to the next most probable state according to
        the transition matrix, and so on.

        Parameters
        ----------
        start_state : int, optional
            Initial state to begin from. Defaults to the most probable
            a priori state.

        Returns
        -------
        new_order : list
            List of states in transmat order.
        """

        # unless specified, start in the a priori most probable state
        if start_state is None:
            start_state = np.argmax(self.startprob_)

        new_order = [start_state]
        num_states = self.transmat_.shape[0]
        rem_states = np.arange(0, start_state).tolist()
        rem_states.extend(np.arange(start_state + 1, num_states).tolist())
        cs = start_state  # current state

        for ii in np.arange(0, num_states - 1):
            # find largest transition to set of remaining states
            nstilde = np.argmax(self.transmat_[cs, rem_states])
            ns = rem_states[nstilde]
            # remove selected state from list of remaining states
            rem_states.remove(ns)
            cs = ns
            new_order.append(cs)

        return new_order

    @property
    def unit_ids(self):
        """
        List of unit IDs associated with the model.

        Returns
        -------
        list
            List of unit IDs.
        """
        return self._unit_ids

    @property
    def unit_labels(self):
        """
        List of unit labels associated with the model.

        Returns
        -------
        list
            List of unit labels.
        """
        return self._unit_labels

    @property
    def means(self):
        """
        Observation matrix (mean firing rates for each state and unit).

        Returns
        -------
        np.ndarray
            Array of shape (n_components, n_units) containing the mean parameters for each state.
        """
        return self.means_

    @property
    def transmat(self):
        """
        Transition probability matrix.

        Returns
        -------
        np.ndarray
            Array of shape (n_components, n_components) where A[i, j] = Pr(S_{t+1}=j | S_t=i).
        """
        return self.transmat_

    @property
    def startprob(self):
        """
        Prior distribution over states.

        Returns
        -------
        np.ndarray
            Array of shape (n_components,) representing the initial state probabilities.
        """
        return self.startprob_

    def get_state_order(self, method=None, start_state=None):
        """
        Return a state ordering, optionally using augmented data.

        Parameters
        ----------
        method : {'transmat', 'mode', 'mean'}, optional
            Method to use for ordering states. 'transmat' (default) uses the transition matrix.
            'mode' or 'mean' use the external mapping (requires self._extern_).
        start_state : int, optional
            Initial state to begin from (used only if method is 'transmat').

        Returns
        -------
        neworder : list
            List of state indices in the new order.

        Notes
        -----
        Both 'mode' and 'mean' assume that _extern_ is in sorted order; this is not verified explicitly.

        Examples
        --------
        >>> order = hmm.get_state_order(method="transmat")
        >>> order = hmm.get_state_order(method="mode")
        """
        if method is None:
            method = "transmat"

        neworder = []

        if method == "transmat":
            return self._get_order_from_transmat(start_state=start_state)
        elif method == "mode":
            if self._extern_ is not None:
                neworder = self._extern_.argmax(axis=1).argsort()
            else:
                raise Exception(
                    "External mapping does not exist yet.First use PoissonHMM.fit_ext()"
                )
        elif method == "mean":
            if self._extern_ is not None:
                (
                    np.tile(np.arange(self._extern_.shape[1]), (self.n_components, 1))
                    * self._extern_
                ).sum(axis=1).argsort()
                neworder = self._extern_.argmax(axis=1).argsort()
            else:
                raise Exception(
                    "External mapping does not exist yet.First use PoissonHMM.fit_ext()"
                )
        else:
            raise NotImplementedError(
                "ordering method '" + str(method) + "' not supported!"
            )
        return neworder

    def _reorder_units_by_ids(self, neworder):
        """
        Reorder unit_ids to match that of a BinnedSpikeTrain.

        WARNING! Modifies self.means_ in-place.

        Parameters
        ----------
        neworder : list or array-like
            List of unit IDs specifying the new order. Must be of size (n_units,).

        Returns
        -------
        self : PoissonHMM
            The reordered PoissonHMM instance.

        Examples
        --------
        >>> hmm._reorder_units_by_ids([3, 1, 2, 0])
        """
        neworder = [self.unit_ids.index(x) for x in neworder]

        oldorder = list(range(len(neworder)))
        for oi, ni in enumerate(neworder):
            frm = oldorder.index(ni)
            to = oi
            swap_cols(self.means_, frm, to)
            self._unit_ids[frm], self._unit_ids[to] = (
                self._unit_ids[to],
                self._unit_ids[frm],
            )
            self._unit_labels[frm], self._unit_labels[to] = (
                self._unit_labels[to],
                self._unit_labels[frm],
            )
            # TODO: re-build unit tags (tag system not yet implemented)
            oldorder[frm], oldorder[to] = oldorder[to], oldorder[frm]

        return self

    def reorder_states(self, neworder):
        """
        Reorder internal HMM states according to a specified order.

        Parameters
        ----------
        neworder : list or array-like
            List of state indices specifying the new order. Must be of size (n_components,).

        Examples
        --------
        >>> hmm.reorder_states([2, 0, 1])
        """
        oldorder = list(range(len(neworder)))
        for oi, ni in enumerate(neworder):
            frm = oldorder.index(ni)
            to = oi
            swap_cols(self.transmat_, frm, to)
            swap_rows(self.transmat_, frm, to)
            swap_rows(self.means_, frm, to)
            if self._extern_ is not None:
                swap_rows(self._extern_, frm, to)
            self.startprob_[frm], self.startprob_[to] = (
                self.startprob_[to],
                self.startprob_[frm],
            )
            oldorder[frm], oldorder[to] = oldorder[to], oldorder[frm]

    def assume_attributes(self, binnedSpikeTrainArray):
        """
        Assume subset of attributes from a BinnedSpikeTrainArray.

        This is used primarily to enable the sampling of sequences after a model has been fit.

        Parameters
        ----------
        binnedSpikeTrainArray : BinnedSpikeTrainArray
            The BinnedSpikeTrainArray instance from which to copy attributes.
        """
        if self._ds is not None:
            warn("PoissonHMM(BinnedSpikeTrain) attributes already exist.")
        for attrib in self.__attributes__:
            exec("self." + attrib + " = binnedSpikeTrainArray." + attrib)
        self._unit_ids = copy.copy(binnedSpikeTrainArray.unit_ids)
        self._unit_labels = copy.copy(binnedSpikeTrainArray.unit_labels)
        self._unit_tags = copy.copy(binnedSpikeTrainArray.unit_tags)

    def _has_same_unit_id_order(self, unit_ids):
        """
        Check if the provided unit_ids are in the same order as the model's unit_ids.

        Parameters
        ----------
        unit_ids : list or array-like
            List of unit IDs to compare.

        Returns
        -------
        bool
            True if the unit_ids are in the same order, False otherwise.

        Raises
        ------
        TypeError
            If the number of unit_ids does not match.
        """
        if self._unit_ids is None:
            return True
        if len(unit_ids) != len(self.unit_ids):
            raise TypeError("Incorrect number of unit_ids encountered!")
        for ii, unit_id in enumerate(unit_ids):
            if unit_id != self.unit_ids[ii]:
                return False
        return True

    def _sliding_window_array(self, bst, w=1):
        """
        Returns an unwrapped data array by sliding w bins one bin at a time.

        If w==1, then bins are non-overlapping.

        Parameters
        ----------
        bst : BinnedSpikeTrainArray
            Input with data array of shape (n_units, n_bins).
        w : int, optional
            Window size (number of bins). Default is 1.

        Returns
        -------
        unwrapped : np.ndarray
            New data array of shape (n_sliding_bins, n_units).
        lengths : np.ndarray
            Array of shape (n_sliding_bins,) indicating the lengths of each window.

        Raises
        ------
        NotImplementedError
            If bst is not a BinnedSpikeTrainArray.
        AssertionError
            If w is not a positive integer.

        Examples
        --------
        >>> unwrapped, lengths = hmm._sliding_window_array(bst, w=3)
        """
        if w is None:
            w = 1
        assert float(w).is_integer(), "w must be a positive integer!"
        assert w > 0, "w must be a positive integer!"

        if not isinstance(bst, BinnedSpikeTrainArray):
            raise NotImplementedError(
                "support for other datatypes not yet implemented!"
            )

        # potentially re-organize internal observation matrix to be
        # compatible with BinnedSpikeTrainArray
        if not self._has_same_unit_id_order(bst.unit_ids):
            self._reorder_units_by_ids(bst.unit_ids)

        if w == 1:
            return bst.data.T, bst.lengths

        n_units, t_bins = bst.data.shape

        # if we decode using multiple bins at a time (w>1) then we have to decode each epoch separately:

        # first, we determine the number of bins we will decode. This requires us to scan over the epochs
        n_bins = 0
        cumlengths = np.cumsum(bst.lengths)
        lengths = np.zeros(bst.n_epochs, dtype=np.int)
        prev_idx = 0
        for ii, to_idx in enumerate(cumlengths):
            datalen = to_idx - prev_idx
            prev_idx = to_idx
            lengths[ii] = np.max((1, datalen - w + 1))

        n_bins = lengths.sum()

        unwrapped = np.zeros((n_units, n_bins))

        # next, we decode each epoch separately, one bin at a time
        cum_lengths = np.insert(np.cumsum(lengths), 0, 0)

        prev_idx = 0
        for ii, to_idx in enumerate(cumlengths):
            data = bst.data[:, prev_idx:to_idx]
            prev_idx = to_idx
            datacum = np.cumsum(
                data, axis=1
            )  # ii'th data segment, with column of zeros prepended
            datacum = np.hstack((np.zeros((n_units, 1)), datacum))
            re = w  # right edge ptr
            # TODO: check if datalen < w and act appropriately
            if lengths[ii] > 1:  # more than one full window fits into data length
                for tt in range(lengths[ii]):
                    obs = (
                        datacum[:, re] - datacum[:, re - w]
                    )  # spikes in window of size w
                    re += 1
                    post_idx = lengths[ii] + tt
                    unwrapped[:, post_idx] = obs
            else:  # only one window can fit in, and perhaps only partially. We just take all the data we can get,
                # and ignore the scaling problem where the window size is now possibly less than bst.ds*w
                post_idx = cum_lengths[ii]
                obs = datacum[:, -1]  # spikes in window of size at most w
                unwrapped[:, post_idx] = obs

        return unwrapped.T, lengths

    def decode(self, X, lengths=None, w=None, algorithm=None):
        """
        Find the most likely state sequence corresponding to ``X``.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features) or BinnedSpikeTrainArray
            Feature matrix of individual samples, or a nelpy.BinnedSpikeTrainArray.
            WARNING: Each decoding window is assumed to be similar in size to those used during training.
            If not, the tuning curves have to be scaled appropriately!
        lengths : array-like of int, shape (n_sequences,), optional
            Lengths of the individual sequences in ``X``. The sum of these should be ``n_samples``.
            Not used when X is a BinnedSpikeTrainArray, in which case the lengths are automatically inferred.
        w : int, optional
            Window size for sliding window decoding (only used for BinnedSpikeTrainArray input).
        algorithm : str, optional
            Decoder algorithm to be used (see DECODER_ALGORITHMS).

        Returns
        -------
        logprob : float or list of float
            Log probability of the produced state sequence.
        state_sequence : np.ndarray or list of np.ndarray
            Labels for each sample from ``X`` obtained via the given decoder algorithm.
        centers : np.ndarray or list of np.ndarray
            Time-centers of all bins contained in ``X``.

        See Also
        --------
        score_samples : Compute the log probability under the model and posteriors.
        score : Compute the log probability under the model.

        Examples
        --------
        >>> logprob, state_seq, centers = hmm.decode(bst)
        >>> logprob, state_seq, centers = hmm.decode(X, algorithm="viterbi")
        """
        if not isinstance(X, BinnedSpikeTrainArray):
            # assume we have a feature matrix
            if w is not None:
                raise NotImplementedError(
                    "sliding window decoding for feature matrices not yet implemented!"
                )
            return self._decode(self, X=X, lengths=lengths, algorithm=algorithm), None
        else:
            # we have a BinnedSpikeTrainArray
            logprobs = []
            state_sequences = []
            centers = []
            for seq in X:
                windowed_arr, lengths = self._sliding_window_array(bst=seq, w=w)
                logprob, state_sequence = self._decode(
                    self, windowed_arr, lengths=lengths, algorithm=algorithm
                )
                logprobs.append(logprob)
                state_sequences.append(state_sequence)
                centers.append(seq.centers)
            return logprobs, state_sequences, centers

    def _decode_from_lambda_only(self, X, lengths=None):
        """
        Decode using the observation (lambda) matrix only (i.e., pure memoryless decoding).

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features) or BinnedSpikeTrainArray
            Feature matrix of individual samples, or a nelpy.BinnedSpikeTrainArray.
            WARNING: Each decoding window is assumed to be similar in size to those used during training.
            If not, the tuning curves have to be scaled appropriately!
        lengths : array-like of int, shape (n_sequences,), optional
            Lengths of the individual sequences in ``X``. The sum of these should be ``n_samples``.
            Not used when X is a BinnedSpikeTrainArray, in which case the lengths are automatically inferred.

        Returns
        -------
        posteriors : list of np.ndarray
            State-membership probabilities for each sample in ``X``; one array for each sequence in X.
        state_sequences : list of np.ndarray
            Labels for each sample from ``X``; one array for each sequence in X.

        Examples
        --------
        >>> posteriors, state_sequences = hmm._decode_from_lambda_only(bst)
        """
        if not isinstance(X, BinnedSpikeTrainArray):
            # assume we have a feature matrix
            raise NotImplementedError("Not yet implemented!")
        else:
            # we have a BinnedSpikeTrainArray
            ratemap = copy.deepcopy(self.means_.T)
            # make sure X and ratemap have same unit_id ordering!
            neworder = [self.unit_ids.index(x) for x in X.unit_ids]
            oldorder = list(range(len(neworder)))
            for oi, ni in enumerate(neworder):
                frm = oldorder.index(ni)
                to = oi
                swap_rows(ratemap, frm, to)
                oldorder[frm], oldorder[to] = oldorder[to], oldorder[frm]

            posteriors = []
            state_sequences = []
            for seq in X:
                posteriors_, cumlengths, mode_pth, mean_pth = decode1D(
                    bst=seq, ratemap=ratemap
                )
                # nanlocs = np.argwhere(np.isnan(mode_pth))
                # state_sequences_ = mode_pth.astype(int)
                state_sequences_ = mode_pth
                posteriors.append(posteriors_)
                state_sequences.append(state_sequences_)

            return posteriors, state_sequences

    def predict_proba(self, X, lengths=None, w=None, returnLengths=False):
        """
        Compute the posterior probability for each state in the model.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features) or BinnedSpikeTrainArray
            Feature matrix of individual samples, or a nelpy.BinnedSpikeTrainArray.
        lengths : array-like of int, shape (n_sequences,), optional
            Lengths of the individual sequences in ``X``. The sum of these should be ``n_samples``.
            Not used when X is a BinnedSpikeTrainArray, in which case the lengths are automatically inferred.
        w : int, optional
            Window size for sliding window decoding (only used for BinnedSpikeTrainArray input).
        returnLengths : bool, optional
            If True, also return the lengths array.

        Returns
        -------
        posteriors : np.ndarray
            Array of shape (n_components, n_samples) with state-membership probabilities for each sample from ``X``.
        lengths : np.ndarray, optional
            Returned if returnLengths is True; array of sequence lengths.

        Examples
        --------
        >>> posteriors = hmm.predict_proba(bst)
        >>> posteriors, lengths = hmm.predict_proba(bst, returnLengths=True)
        """
        if not isinstance(X, BinnedSpikeTrainArray):
            print("we have a " + str(type(X)))
            # assume we have a feature matrix
            if w is not None:
                raise NotImplementedError(
                    "sliding window decoding for feature matrices not yet implemented!"
                )
            if returnLengths:
                return np.transpose(
                    self._predict_proba(self, X, lengths=lengths)
                ), lengths
            return np.transpose(self._predict_proba(self, X, lengths=lengths))
        else:
            # we have a BinnedSpikeTrainArray
            windowed_arr, lengths = self._sliding_window_array(bst=X, w=w)
            if returnLengths:
                return np.transpose(
                    self._predict_proba(self, windowed_arr, lengths=lengths)
                ), lengths
            return np.transpose(
                self._predict_proba(self, windowed_arr, lengths=lengths)
            )

    def predict(self, X, lengths=None, w=None):
        """
        Find the most likely state sequence corresponding to ``X``.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features) or BinnedSpikeTrainArray
            Feature matrix of individual samples, or a nelpy.BinnedSpikeTrainArray.
        lengths : array-like of int, shape (n_sequences,), optional
            Lengths of the individual sequences in ``X``. The sum of these should be ``n_samples``.
            Not used when X is a BinnedSpikeTrainArray, in which case the lengths are automatically inferred.
        w : int, optional
            Window size for sliding window decoding (only used for BinnedSpikeTrainArray input).

        Returns
        -------
        state_sequence : np.ndarray or list of np.ndarray
            Labels for each sample from ``X``.

        Examples
        --------
        >>> state_seq = hmm.predict(bst)
        >>> state_seq = hmm.predict(X)
        """
        _, state_sequences, centers = self.decode(X=X, lengths=lengths, w=w)
        return state_sequences

    def sample(self, n_samples=1, random_state=None):
        """
        Generate random samples from the model.

        Parameters
        ----------
        n_samples : int
            Number of samples to generate.
        random_state : RandomState or int, optional
            A random number generator instance or seed. If None, the object's random_state is used.

        Returns
        -------
        X : np.ndarray
            Feature matrix of shape (n_samples, n_features).
        state_sequence : np.ndarray
            State sequence produced by the model.

        Examples
        --------
        >>> X, states = hmm.sample(n_samples=100)
        """
        return self._sample(self, n_samples=n_samples, random_state=random_state)

    def score_samples(self, X, lengths=None, w=None):
        """Compute the log probability under the model and compute posteriors.

        Parameters
        ----------
        X : array-like, shape (n_samples, n_features)
            Feature matrix of individual samples.
            OR
            nelpy.BinnedSpikeTrainArray
        lengths : array-like of integers, shape (n_sequences, ), optional
            Lengths of the individual sequences in ``X``. The sum of
            these should be ``n_samples``. This is not used when X is
            a nelpy.BinnedSpikeTrainArray, in which case the lenghts are
            automatically inferred.

        Returns
        -------
        logprob : float
            Log likelihood of ``X``; one scalar for each sequence in X.

        posteriors : array, shape (n_components, n_samples)
            State-membership probabilities for each sample in ``X``;
            one array for each sequence in X.

        See Also
        --------
        score : Compute the log probability under the model.
        decode : Find most likely state sequence corresponding to ``X``.
        """

        if not isinstance(X, BinnedSpikeTrainArray):
            # assume we have a feature matrix
            if w is not None:
                raise NotImplementedError(
                    "sliding window decoding for feature matrices not yet implemented!"
                )
            logprobs, posteriors = self._score_samples(self, X, lengths=lengths)
            return (
                logprobs,
                posteriors,
            )  # .T why does this transpose affect hmm.predict_proba!!!????
        else:
            # we have a BinnedSpikeTrainArray
            logprobs = []
            posteriors = []
            for seq in X:
                windowed_arr, lengths = self._sliding_window_array(bst=seq, w=w)
                logprob, posterior = self._score_samples(
                    self, X=windowed_arr, lengths=lengths
                )
                logprobs.append(logprob)
                posteriors.append(posterior.T)
            return logprobs, posteriors

    def score(self, X, lengths=None, w=None):
        """Compute the log probability under the model.

        Parameters
        ----------
        X : array-like, shape (n_samples, n_features)
            Feature matrix of individual samples.
            OR
            nelpy.BinnedSpikeTrainArray
        lengths : array-like of integers, shape (n_sequences, ), optional
            Lengths of the individual sequences in ``X``. The sum of
            these should be ``n_samples``. This is not used when X is
            a nelpy.BinnedSpikeTrainArray, in which case the lenghts are
            automatically inferred.

        Returns
        -------
        logprob : float, or list of floats
            Log likelihood of ``X``; one scalar for each sequence in X.

        See Also
        --------
        score_samples : Compute the log probability under the model and
            posteriors.
        decode : Find most likely state sequence corresponding to ``X``.
        """

        if not isinstance(X, BinnedSpikeTrainArray):
            # assume we have a feature matrix
            if w is not None:
                raise NotImplementedError(
                    "sliding window decoding for feature matrices not yet implemented!"
                )
            return self._score(self, X, lengths=lengths)
        else:
            # we have a BinnedSpikeTrainArray
            logprobs = []
            for seq in X:
                windowed_arr, lengths = self._sliding_window_array(bst=seq, w=w)
                logprob = self._score(self, X=windowed_arr, lengths=lengths)
                logprobs.append(logprob)
        return logprobs

    def _cum_score_per_bin(self, X, lengths=None, w=None):
        """Compute the log probability under the model, cumulatively for each bin per event."""

        if not isinstance(X, BinnedSpikeTrainArray):
            # assume we have a feature matrix
            if w is not None:
                raise NotImplementedError(
                    "sliding window decoding for feature matrices not yet implemented!"
                )
            return self._score(self, X, lengths=lengths)
        else:
            # we have a BinnedSpikeTrainArray
            logprobs = []
            for seq in X:
                windowed_arr, lengths = self._sliding_window_array(bst=seq, w=w)
                n_bins, _ = windowed_arr.shape
                for ii in range(1, n_bins + 1):
                    logprob = self._score(self, X=windowed_arr[:ii, :])
                    logprobs.append(logprob)
        return logprobs

    def fit(self, X, lengths=None, w=None):
        """Estimate model parameters using nelpy objects.

        An initialization step is performed before entering the
        EM-algorithm. If you want to avoid this step for a subset of
        the parameters, pass proper ``init_params`` keyword argument
        to estimator's constructor.

        Parameters
        ----------
        X : array-like, shape (n_samples, n_units)
            Feature matrix of individual samples.
            OR
            nelpy.BinnedSpikeTrainArray
        lengths : array-like of integers, shape (n_sequences, )
            Lengths of the individual sequences in ``X``. The sum of
            these should be ``n_samples``. This is not used when X is
            a nelpy.BinnedSpikeTrainArray, in which case the lenghts are
            automatically inferred.

        Returns
        -------
        self : object
            Returns self.
        """
        if not isinstance(X, BinnedSpikeTrainArray):
            # assume we have a feature matrix
            if w is not None:
                raise NotImplementedError(
                    "sliding window decoding for feature matrices not yet implemented!"
                )
            self._fit(self, X, lengths=lengths)
        else:
            # we have a BinnedSpikeTrainArray
            windowed_arr, lengths = self._sliding_window_array(bst=X, w=w)
            self._fit(self, windowed_arr, lengths=lengths)
            # adopt unit_ids, unit_labels, etc. from BinnedSpikeTrain
            self.assume_attributes(X)
        return self

    def fit_ext(
        self,
        X,
        ext,
        n_extern=None,
        lengths=None,
        save=True,
        w=None,
        normalize=True,
        normalize_by_occupancy=True,
    ):
        """Learn a mapping from the internal state space, to an external
        augmented space (e.g. position).

        Returns a row-normalized version of (n_states, n_ext), that
        is, a distribution over external bins for each state.

        X : BinnedSpikeTrainArray

        ext : array-like
            array of external correlates (n_bins, )
        n_extern : int
            number of extern variables, with range 0,.. n_extern-1
        save : bool
            stores extern in PoissonHMM if true, discards it if not
        w:
        normalize : bool
            If True, then normalize each state to have a distribution over ext.
        occupancy : array of bin counts
            Default is all ones (uniform).

        self.extern_ of size (n_components, n_extern)
        """

        if n_extern is None:
            n_extern = len(unique(ext))
            ext_map = np.arange(n_extern)
            for ii, ele in enumerate(unique(ext)):
                ext_map[ele] = ii
        else:
            ext_map = np.arange(n_extern)

        # idea: here, ext can be anything, and n_extern should be range
        # we can e.g., define extern correlates {leftrun, rightrun} and
        # fit the mapping. This is not expected to be good at all for
        # most states, but it could allow us to identify a state or two
        # for which there *might* be a strong predictive relationship.
        # In this way, the binning, etc. should be done external to this
        # function, but it might still make sense to encapsulate it as
        # a helper function inside PoissonHMM?

        # xpos, ypos = get_position(exp_data['session1']['posdf'], bst.centers)
        # x0=0; xl=100; n_extern=50
        # xx_left = np.linspace(x0,xl,n_extern+1)
        # xx_mid = np.linspace(x0,xl,n_extern+1)[:-1]; xx_mid += (xx_mid[1]-xx_mid[0])/2
        # ext = np.digitize(xpos, xx_left) - 1 # spatial bin numbers

        extern = np.zeros((self.n_components, n_extern))

        if not isinstance(X, BinnedSpikeTrainArray):
            # assume we have a feature matrix
            if w is not None:
                raise NotImplementedError(
                    "sliding window decoding for feature matrices not yet implemented!"
                )
            posteriors = self.predict_proba(X=X, lengths=lengths, w=w)
        else:
            # we have a BinnedSpikeTrainArray
            posteriors = self.predict_proba(X=X, lengths=lengths, w=w)

        posteriors = np.vstack(posteriors.T)  # 1D array of states, of length n_bins

        if len(posteriors) != len(ext):
            raise ValueError("ext must have same length as decoded state sequence!")

        for ii, posterior in enumerate(posteriors):
            if not np.isnan(ext[ii]):
                extern[:, ext_map[int(ext[ii])]] += np.transpose(posterior)

        if normalize_by_occupancy:
            occupancy, _ = np.histogram(ext, bins=n_extern, range=[0, n_extern])
            occupancy[occupancy == 0] = 1
            occupancy = np.atleast_2d(occupancy)
        else:
            occupancy = 1

        extern = extern / occupancy

        if normalize:
            # normalize extern tuning curves:
            rowsum = np.tile(extern.sum(axis=1), (n_extern, 1)).T
            rowsum = np.where(np.isclose(rowsum, 0), 1, rowsum)
            extern = extern / rowsum

        if save:
            self._extern_ = extern
            # self._extern_map = ext_map

        return extern

    def fit_ext2(self, X, ext, n_extern=None, lengths=None, w=None):
        """Learn a mapping from the internal state space, to an external
        augmented space (e.g. position).

        Returns a column-normalized version of (n_states, n_ext), that
        is, a distribution over states for each extern bin.

        X : BinnedSpikeTrainArray

        ext : array-like
            array of external correlates (n_bins, )
        n_extern : int
            number of extern variables, with range 0,.. n_extern-1

        save : bool
            stores extern in PoissonHMM if true, discards it if not

        self.extern_ of size (n_components, n_extern)
        """

        ext_map = np.arange(n_extern)
        if n_extern is None:
            n_extern = len(unique(ext))
            for ii, ele in enumerate(unique(ext)):
                ext_map[ele] = ii

        # idea: here, ext can be anything, and n_extern should be range
        # we can e.g., define extern correlates {leftrun, rightrun} and
        # fit the mapping. This is not expexted to be good at all for
        # most states, but it could allow us to identify a state or two
        # for which there *might* be a strong predictive relationship.
        # In this way, the binning, etc. should be done external to this
        # function, but it might still make sense to encapsulate it as
        # a helper function inside PoissonHMM?

        # xpos, ypos = get_position(exp_data['session1']['posdf'], bst.centers)
        # x0=0; xl=100; n_extern=50
        # xx_left = np.linspace(x0,xl,n_extern+1)
        # xx_mid = np.linspace(x0,xl,n_extern+1)[:-1]; xx_mid += (xx_mid[1]-xx_mid[0])/2
        # ext = np.digitize(xpos, xx_left) - 1 # spatial bin numbers

        extern = np.zeros((self.n_components, n_extern))

        if not isinstance(X, BinnedSpikeTrainArray):
            # assume we have a feature matrix
            if w is not None:
                raise NotImplementedError(
                    "sliding window decoding for feature matrices not yet implemented!"
                )
            posteriors = self.predict_proba(X=X, lengths=lengths, w=w)
        else:
            # we have a BinnedSpikeTrainArray
            posteriors = self.predict_proba(X=X, lengths=lengths, w=w)
        posteriors = np.vstack(posteriors.T)  # 1D array of states, of length n_bins

        if len(posteriors) != len(ext):
            raise ValueError("ext must have same length as decoded state sequence!")

        for ii, posterior in enumerate(posteriors):
            if not np.isnan(ext[ii]):
                extern[:, ext_map[int(ext[ii])]] += np.transpose(posterior)

        # normalize extern tuning curves:
        colsum = np.tile(extern.sum(axis=0), (self.n_components, 1))
        colsum = np.where(np.isclose(colsum, 0), 1, colsum)
        extern = extern / colsum

        return extern

    def decode_ext(self, X, lengths=None, w=None, ext_shape=None):
        """
        Find memoryless most likely state sequence corresponding to ``X``,
        (that is, the symbol-by-symbol MAP sequence) and then map those
        states to an associated external representation (e.g., position).

        Parameters
        ----------
        X : array-like, shape (n_samples, n_features) or BinnedSpikeTrainArray
            Feature matrix of individual samples, or a nelpy.BinnedSpikeTrainArray.
        lengths : array-like of integers, shape (n_sequences, ), optional
            Lengths of the individual sequences in ``X``. The sum of these should be ``n_samples``.
            Not used when X is a nelpy.BinnedSpikeTrainArray, in which case the lengths are automatically inferred.
        w : int, optional
            Window size for sliding window decoding (only used for BinnedSpikeTrainArray input).
        ext_shape : tuple, optional
            Shape of the external variables.

        Returns
        -------
        ext_posteriors : np.ndarray
            Array of shape (n_extern, n_samples) with state-membership probabilities for each sample in ``X``.
        bdries : np.ndarray
            Array of bin boundaries.
        mode_pth : np.ndarray
            Most likely external variable sequence (mode path).
        mean_pth : np.ndarray
            Mean external variable sequence (mean path).

        Examples
        --------
        For 1D external variables:
        >>> posterior_pos, bdries, mode_pth, mean_pth = hmm.decode_ext(
        ...     bst_no_ripple, ext_shape=(vtc.n_bins,)
        ... )
        >>> mean_pth = vtc.bins[0] + mean_pth * (vtc.bins[-1] - vtc.bins[0])

        For 2D external variables:
        >>> posterior_, bdries_, mode_pth_, mean_pth_ = hmm.decode_ext(
        ...     bst, ext_shape=(ext_nx, ext_ny)
        ... )
        >>> mean_pth_[0, :] = vtc2d.xbins[0] + mean_pth_[0, :] * (
        ...     vtc2d.xbins[-1] - vtc2d.xbins[0]
        ... )
        >>> mean_pth_[1, :] = vtc2d.ybins[0] + mean_pth_[1, :] * (
        ...     vtc2d.ybins[-1] - vtc2d.ybins[0]
        ... )
        """

        _, n_extern = self._extern_.shape

        if ext_shape is None:
            ext_shape = n_extern

        if not isinstance(X, BinnedSpikeTrainArray):
            # assume we have a feature matrix
            raise NotImplementedError("not implemented yet.")
            if w is not None:
                raise NotImplementedError(
                    "sliding window decoding for feature matrices not yet implemented!"
                )
        else:
            # we have a BinnedSpikeTrainArray
            pass
        if len(ext_shape) == 1:
            # do old style decoding
            # TODO: this can be improved to be like the 2D case!
            state_posteriors, lengths = self.predict_proba(
                X=X, lengths=lengths, w=w, returnLengths=True
            )
            # fixy = np.mean(self._extern_ * np.arange(n_extern), axis=1)
            # mean_pth = np.sum(state_posteriors.T*fixy, axis=1) # range 0 to 1
            ext_posteriors = np.dot(
                (self._extern_ * np.arange(n_extern)).T, state_posteriors
            )
            # normalize ext_posterior distributions:
            ext_posteriors = ext_posteriors / ext_posteriors.sum(axis=0)
            mean_pth = (
                ext_posteriors.T * np.atleast_2d(np.linspace(0, 1, n_extern))
            ).sum(axis=1)
            mode_pth = (
                np.argmax(ext_posteriors, axis=0) / n_extern
            )  # range 0 to n_extern

        elif len(ext_shape) == 2:
            ext_posteriors = np.zeros((ext_shape[0], ext_shape[1], X.n_bins))
            # get posterior distribution over states, of size (num_States, n_extern)
            state_posteriors, lengths = self.predict_proba(
                X=X, lengths=lengths, w=w, returnLengths=True
            )
            # for each bin, compute the distribution in the external domain
            for bb in range(X.n_bins):
                ext_posteriors[:, :, bb] = np.reshape(
                    (self._extern_ * state_posteriors[:, [bb]]).sum(axis=0), ext_shape
                )
            # now compute mean and mode paths
            expected_x = np.sum(
                (
                    ext_posteriors.sum(axis=1)
                    * np.atleast_2d(np.linspace(0, 1, ext_shape[0])).T
                ),
                axis=0,
            )
            expected_y = np.sum(
                (
                    ext_posteriors.sum(axis=0)
                    * np.atleast_2d(np.linspace(0, 1, ext_shape[1])).T
                ),
                axis=0,
            )
            mean_pth = np.vstack((expected_x, expected_y))

            mode_pth = np.zeros((2, X.n_bins))
            for tt in range(X.n_bins):
                if np.any(np.isnan(ext_posteriors[:, :, tt])):
                    mode_pth[0, tt] = np.nan
                    mode_pth[0, tt] = np.nan
                else:
                    x_, y_ = np.unravel_index(
                        np.argmax(ext_posteriors[:, :, tt]),
                        (ext_shape[0], ext_shape[1]),
                    )
                    mode_pth[0, tt] = x_ / ext_shape[0]
                    mode_pth[1, tt] = y_ / ext_shape[1]

            ext_posteriors = np.transpose(ext_posteriors, axes=[1, 0, 2])
        else:
            raise TypeError("shape not currently supported!")

        bdries = np.cumsum(lengths)

        return ext_posteriors, bdries, mode_pth, mean_pth

    def _plot_external(
        self,
        *,
        figsize=(3, 5),
        sharey=True,
        labelstates=None,
        ec=None,
        fillcolor=None,
        lw=None,
    ):
        """plot the externally associated state<-->extern mapping

        WARNING! This function is not complete, and hence 'private',
        and may be moved somewhere else later on.
        """

        if labelstates is None:
            labelstates = [1, self.n_components]
        if ec is None:
            ec = "k"
        if fillcolor is None:
            fillcolor = "gray"
        if lw is None:
            lw = 1.5

        fig, axes = subplots(self.n_components, 1, figsize=figsize, sharey=sharey)

        xvals = np.arange(len(self._extern_.T[:, 0]))

        for state, ax in enumerate(axes):
            ax.fill_between(xvals, 0, self._extern_.T[:, state], color=fillcolor)
            ax.plot(xvals, self._extern_.T[:, state], color=ec, lw=lw)
            if state + 1 in labelstates:
                ax.set_ylabel(str(state + 1), rotation=0, y=-0.1)
            ax.set_xticklabels([])
            ax.set_yticklabels([])
            ax.spines["right"].set_visible(False)
            ax.spines["top"].set_visible(False)
            ax.spines["bottom"].set_visible(False)
            ax.spines["left"].set_visible(False)
            plotting.utils.no_yticks(ax)
            plotting.utils.no_xticks(ax)
        # fig.suptitle('normalized place fields sorted by peak location (left) and mean location (right)', y=0.92, fontsize=14)
        # ax.set_xticklabels(['0','20', '40', '60', '80', '100'])
        ax.set_xlabel("external variable")
        fig.text(
            0.02, 0.5, "normalized state distribution", va="center", rotation="vertical"
        )

        return fig, ax

    def estimate_model_quality(self, bst, *, n_shuffles=1000, k_folds=5, verbose=False):
        """Estimate the HMM 'model quality' associated with the set of events in bst.

        TODO: finish docstring, and do some more consistency checking...

        Params
        ======

        Returns
        =======

        quality :
        scores :
        shuffled :

        """
        n_states = self.n_components
        quality, scores, shuffles = estimate_model_quality(
            bst=bst,
            n_states=n_states,
            n_shuffles=n_shuffles,
            k_folds=k_folds,
            verbose=False,
        )

        return quality, scores, shuffles