@@ -13,117 +13,28 @@ This way, we can quantify phenotypic heterogeneity within a population of cells
 We first introduce how to synthesize these populations by working up from the basic unit; cells. We then introduce how to synthesize lineages, which are just hierarchical lineage tree groupings of cells based on their family history. Our model ultimately analyzes populations which are groups of one or more lineages that share the same states.
 
 
-1. Synthesizing Cells
----------------------
+1. Lineage Tree Data Model
+---------------------------
 
-Users will rarely ever have to create individual cells, but their
-functionality is presented so the user is familiar with them. Users will
-however, create lineages, and as such, they will have to become
-comfortable with creating transition matrices using ``numpy``.
-Transition matrices define the probabilities of how cells will divide.
-Cells can only be made (new cells as a result of an existing cell
-dividing) when given a transition matrix. We provide an example below.
+Lineage trees in tHMM are represented at the lineage level using a **SciPy CSR sparse array** (representing parent-to-daughter connectivity) alongside **NumPy arrays** for observations, hidden states, and cell lifetimes.
 
-The transition matrix defines the rate at which cells are likely to
-change from one state to another. We use the defintion of a stochastic
-transition rate matrix from Wikipedia, that is, the column index defines
-the state in which we start, and the row defines the state at which we
-end up. That is, if an element at the index of a transition matrix
-:math:`T` at row :math:`i` and at column :math:`j`, is defined as
-:math:`T_{i,j}`, then
+Transition matrices define the probabilities of how cells divide and switch states across generations.
 
-.. math:: T_{i,j} = \mathbb{P}(z = j | z = i).
+The transition matrix defines the rate at which cells change from one state to another. Specifically, if an element of a transition matrix :math:`T` at row :math:`i` and column :math:`j` is defined as :math:`T_{i,j}`, then
 
+.. math:: T_{i,j} = \mathbb{P}(z_{\text{daughter}} = j | z_{\text{parent}} = i).
 
-\ Indexing for this matrix and the states starts at :math:`0`. Usually
-the number of states is represented as the capital letter :math:`K` and
-indexed by the lower-case letter :math:`k`. For most of our examples, we
-will deal with two states, i.e., :math:`K=2`.
+Indexing for states starts at :math:`0`. Usually the number of states is represented as the capital letter :math:`K` and indexed by :math:`k`. For most examples, we deal with two states, i.e., :math:`K=2`.
 
 .. code:: ipython3
 
     import numpy as np
-    # The capital letter T will represent the state transition probability matrix.
-    # This transition matrix, in particular, is the two-dimensional 
-    # identity matrix. Recall that the elements of the transition matrix
-    # represent the probability of transitioning from one state to another.
-    T = np.array([[1.0, 0.0],
-                  [0.0, 1.0]],  dtype="float")
-    # This identity matrix implies that there are no transitions between cells of
-    # state 0 to cells of state 1. It also implies that there are no transitions 
-    # between cells of state 1 to cells of state 0.
-    # Another transition matrix could be the following:
-    #
-    # T = np.array([[0.75, 0.25],
-    #               [0.15, 0.85]], dtype="float")
-    #
-    # Note that the rows of the transition matrix have to sum one, because of the
-    # Law of Total Probability and the Law of Conditional Probability. This will be
-    # expanded on in a later notebook.
-
-Cells are defined by their state, their relationships to other cells,
-and collections of observations. Knowing how to create cells, however,
-is not required by the user. It is beneficial to understand how the
-``CellVar`` class is designed to create objects that store a ``state``,
-its relationships (``left``, ``right``, ``parent``) to other cells, and
-its multivariate observations (``obs``). When cells are created via the
-``member`` function, their familiar relationships are automatically
-assigned. The two daughters are assigned to ``left`` and ``right`` and
-for the daughters, the parent is assigned to ``parent``. Other familiar
-relationships can be accessed through other member variables and
-functions. Note that the first generation is empty (we can see this
-because the ``parent_cell`` is instantiated with ``parent=None`` and
-with ``gen=1`` which means that ``gen=0`` of the lineage contains
-nothing except ``None``). As such, indexing for generations of lineages
-starts at ``1``. We will discuss more about the multivariate
-observations that cells store later; for now, the following exercises
-will focus on the transition matrix and the familiar relationships of
-cells.
 
-.. code:: ipython3
-
-    import numpy as np
-    from lineage.CellVar import CellVar as c
-
-    parent_cell = c(state=0, parent=None)
+    # State transition probability matrix
     T = np.array([[0.75, 0.25],
                   [0.15, 0.85]], dtype="float")
-    left_cell, right_cell = parent_cell.divide(T)
-
---------------
-
-.. code:: ipython3
-
-    print(parent_cell)
-    # <lineage.CellVar.CellVar object at 0x000001A685C55E10>
-    print(left_cell, right_cell)
-    # <lineage.CellVar.CellVar object at 0x000001A6961B3EF0> <lineage.CellVar.CellVar object at 0x000001A69609C978>
-    print(f"\nThe value of the element at (0,0) of the transition rate matrix is {T[0,0]}")
-    # The value of the element at (0,0) of the transition rate matrix is 0.75
-
---------------
-
-The ``gen`` argument for instantiating cells represents the generation
-of the cell which start at 1.
-
-
---------------
-
-
-``parent_cell``, ``left_cell``, and ``right_cell`` define a 3-cell
-lineage, with 2 generations. The first generation has one cell which was
-declared and can be accessed at ``parent_cell``. Calling the member
-function ``divide`` on ``parent_cell`` created two new cells which can
-be accessed at ``left_cell`` and ``right_cell``. The daughter cells of
-any cell can also be accessed using “dot” notation, using the
-member variables, ``left`` and ``right``. The division process
-utilizes the transition matrix. Our code provides some very basic
-printing methods to print out cells. The following code verifies that the ``left_cell`` is actually the left daughter of ``parent_cell``.
-
-.. code:: ipython3
 
-    # Use the `is` keyword to compare Python objects.
-    assert left_cell is parent_cell.left
+Note that the rows of the transition matrix must sum to 1 by the Law of Total Probability.
 
 --------------
 
 
@@ -28,12 +28,12 @@ def do_E_step(tHMMobj: tHMM) -> tuple[list, list, list, list]:
     EL = get_Emission_Likelihoods(tHMMobj.X, tHMMobj.estimate.E)
 
     for ii, lO in enumerate(tHMMobj.X):
-        MSD.append(get_MSD(lO.cell_to_daughters, tHMMobj.estimate.pi, tHMMobj.estimate.T))
+        MSD.append(get_MSD(lO.tree, tHMMobj.estimate.pi, tHMMobj.estimate.T))
 
-        NF_one, beta = get_beta_and_NF(lO.leaves_idx, lO.cell_to_daughters, tHMMobj.estimate.T, MSD[ii], EL[ii])
+        NF_one, beta = get_beta_and_NF(lO.leaves_idx, lO.tree, tHMMobj.estimate.T, MSD[ii], EL[ii])
         NF.append(NF_one)
         betas.append(beta)
-        gammas.append(get_gamma(lO.cell_to_daughters, tHMMobj.estimate.T, MSD[ii], betas[ii]))
+        gammas.append(get_gamma(lO.tree, tHMMobj.estimate.T, MSD[ii], betas[ii]))
 
     return MSD, NF, betas, gammas
 
@@ -165,8 +165,7 @@ def do_M_T_step(
         for num, lO in enumerate(tt.X):
             # local T estimate
             numer_e += get_all_zetas(
-                lO.leaves_idx,
-                lO.cell_to_daughters,
+                lO.tree,
                 betas[i][num],
                 MSD[i][num],
                 gammas[i][num],
@@ -193,8 +192,7 @@ def do_M_E_step(tHMMobj: tHMM, gammas: list[np.ndarray]):
     :type tHMMobj: object
     :param gammas: gamma values. The conditional probability of states, given the observation of the whole tree
     """
-    all_cells = [cell.obs for lineage in tHMMobj.X for cell in lineage.output_lineage]
-    cell_arr = np.array(all_cells)
+    cell_arr = np.vstack([lineage.obs for lineage in tHMMobj.X])
     all_gammas = np.vstack(gammas)
     for state_j in range(tHMMobj.num_states):
         tHMMobj.estimate.E[state_j].estimator(cell_arr, all_gammas[:, state_j])
@@ -206,24 +204,19 @@ def do_M_E_step_atonce(all_tHMMobj: list[tHMM], all_gammas: list[list[np.ndarray
     After reshaping, we will have a list of lists for each state.
     This function is specifically written for the experimental data of G1 and S-G2 cell cycle fates and durations.
     """
-    gms = []
-    for gm in all_gammas:
-        gms.append(np.vstack(gm))
+    gms = [np.vstack(gm) for gm in all_gammas]
 
-    all_cells = np.array([cell.obs for lineage in all_tHMMobj[0].X for cell in lineage.output_lineage])
-    if len(all_cells[1, :]) == 6:
-        phase = True
-    else:
-        phase = False
+    all_cells_0 = np.vstack([lineage.obs for lineage in all_tHMMobj[0].X])
+    phase = all_cells_0.shape[1] == 6
 
     G1cells = []
     G2cells = []
     cells = []
     for tHMMobj in all_tHMMobj:
-        all_cells = np.array([cell.obs for lineage in tHMMobj.X for cell in lineage.output_lineage])
+        all_cells = np.vstack([lineage.obs for lineage in tHMMobj.X])
         if phase:
-            G1cells.append(all_cells[:, np.array([0, 2, 4])])
-            G2cells.append(all_cells[:, np.array([1, 3, 5])])
+            G1cells.append(all_cells[:, [0, 2, 4]])
+            G2cells.append(all_cells[:, [1, 3, 5]])
         else:
             cells.append(all_cells)
 
 
@@ -1,10 +1,10 @@
 import numpy as np
 import numpy.typing as npt
-from numba import jit
+from scipy.sparse import csr_array
 
 
 def get_MSD(
-    cell_to_daughters: np.ndarray,
+    tree: csr_array,
     pi: npt.NDArray[np.float64],
     T: npt.NDArray[np.float64],
 ) -> npt.NDArray[np.float64]:
@@ -25,27 +25,34 @@ def get_MSD(
 
     :math:`P(z_u = k) = \sum_j(Transition(j -> k) * P(parent_{cell_u}) = j)`
 
+    :param tree: CSR array representing the lineage tree adjacency
     :param pi: Initial probabilities vector
     :param T: State transitions matrix
     :return: The marginal state distribution
     """
-    m = np.zeros((cell_to_daughters.shape[0], pi.size))
+    m = np.zeros((tree.shape[0], pi.size))
     m[0, :] = pi
 
     # recursion based on parent cell
-    for pIDX, cIDX in enumerate(cell_to_daughters):
-        if cIDX[0] != -1:
-            m[cIDX[0], :] = m[pIDX, :] @ T
-        if cIDX[1] != -1:
-            m[cIDX[1], :] = m[pIDX, :] @ T
+    for pIDX in range(tree.shape[0]):
+        start, end = tree.indptr[pIDX], tree.indptr[pIDX + 1]
+        if start < end:
+            p_trans = m[pIDX, :] @ T
+            for cIDX in tree.indices[start:end]:
+                m[cIDX, :] = p_trans
 
     # Assert all ~= 1.0
     assert np.linalg.norm(np.sum(m, axis=1) - 1.0) < 1e-9
     return m
 
 
-@jit
-def get_beta_and_NF(leaves_idx, cell_to_daughters, T: np.ndarray, MSD: np.ndarray, EL: np.ndarray):
+def get_beta_and_NF(
+    leaves_idx: np.ndarray,
+    tree: csr_array,
+    T: np.ndarray,
+    MSD: np.ndarray,
+    EL: np.ndarray,
+) -> tuple[np.ndarray, np.ndarray]:
     r"""
     Normalizing factor (NF) matrix and base case at the leaves.
 
@@ -54,50 +61,9 @@ def get_beta_and_NF(leaves_idx, cell_to_daughters, T: np.ndarray, MSD: np.ndarra
     This normalizing factor is essentially the marginal
     observation distribution for a node.
 
-    This function gets the normalizing factor for
-    the upward recursion only for the leaves.
-    We first calculate the joint probability
-    using the definition of conditional probability:
-
-    :math:`P(x_n = x | z_n = k) * P(z_n = k) = P(x_n = x , z_n = k)`,
-    where n are the leaf nodes.
-
-    We can then sum this joint probability over k,
-    which are the possible states z_n can be,
-    and through the law of total probability,
-    obtain the marginal observation distribution
-    :math:`P(x_n = x) = sum_k ( P(x_n = x , z_n = k) ) = P(x_n = x)`.
-
-    Beta matrix and base case at the leaves.
-
-    Each element in this N by K matrix is the beta value
-    for each cell and at each state. In particular, this
-    value is derived from the Marginal State Distributions
-    (MSD), the Emission Likelihoods (EL), and the
-    Normalizing Factors (NF). Each beta value
-    for the leaves is exactly the probability
-
-    :math:`beta[n,k] = P(z_n = k | x_n = x)`.
-
-    Using Bayes Theorem, we see that the above equals
-
-    numerator = :math:`P(x_n = x | z_n = k) * P(z_n = k)`
-    denominator = :math:`P(x_n = x)`
-    :math:`beta[n,k] = numerator / denominator`
-
-    For non-leaf cells, the first value in the numerator is the Emission
-    Likelihoods. The second value in the numerator is
-    the Marginal State Distributions. The value in the
-    denominator is the Normalizing Factor.
-
-    Traverses upward through each tree and calculates the
-    beta value for each non-leaf cell. The normalizing factors (NFs)
-    are also calculated as an intermediate for determining each
-    beta term. Helper functions are called to determine one of
-    the terms in the NF equation. This term is also used in the calculation
-    of the betas.
-
-    :param tHMMobj: A class object with properties of the lineages of cells
+    :param leaves_idx: array of indices corresponding to leaf cells
+    :param tree: CSR array representing the lineage tree adjacency
+    :param T: Transition probability matrix
     :param MSD: The marginal state distribution P(z_n = k)
     :param EL: The emissions likelihood
     :return: normalizing factor. The marginal observation distribution P(x_n = x)
@@ -108,32 +74,23 @@ def get_beta_and_NF(leaves_idx, cell_to_daughters, T: np.ndarray, MSD: np.ndarra
     ELMSD = EL * MSD
 
     ### NF leaf calculation
-    # P(x_n = x , z_n = k) = P(x_n = x | z_n = k) * P(z_n = k)
-    # this product is the joint probability
-    # P(x_n = x) = sum_k ( P(x_n = x , z_n = k) )
-    # the sum of the joint probabilities is the marginal probability
     NF = np.zeros(MSD.shape[0], dtype=float)  # instantiating N by 1 array
     NF[leaves_idx] = np.sum(ELMSD[leaves_idx, :], axis=1)
     assert np.all(np.isfinite(NF))
 
     ### beta calculation
     beta = np.zeros_like(MSD)
-
-    # Emission Likelihood, Marginal State Distribution, Normalizing Factor (same regardless of state)
-    # P(x_n = x | z_n = k), P(z_n = k), P(x_n = x)
     beta[leaves_idx, :] = ELMSD[leaves_idx, :] / NF[leaves_idx, np.newaxis]
 
     # Assert all ~= 1.0
     assert np.abs(np.sum(beta[-1]) - 1.0) < 1e-9
 
-    cIDXs = np.arange(MSD.shape[0])
-    cIDXs = np.delete(cIDXs, leaves_idx)
-    cIDXs = np.flip(cIDXs)
+    cIDXs = np.nonzero(np.diff(tree.indptr) > 0)[0][::-1]
 
     for pii in cIDXs:
-        ch_ii = cell_to_daughters[pii, :]
+        ch_ii = tree.indices[tree.indptr[pii] : tree.indptr[pii + 1]]
         ratt = (beta[ch_ii, :] / MSD_array[ch_ii, :]) @ T.T
-        fac1 = ratt[0, :] * ratt[1, :] * ELMSD[pii, :]
+        fac1 = np.prod(ratt, axis=0) * ELMSD[pii, :]
 
         NF[pii] = np.sum(fac1)
         beta[pii, :] = fac1 / NF[pii]
@@ -142,7 +99,7 @@ def get_beta_and_NF(leaves_idx, cell_to_daughters, T: np.ndarray, MSD: np.ndarra
 
 
 def get_gamma(
-    cell_to_daughters: npt.NDArray[np.uintp],
+    tree: csr_array,
     T: npt.NDArray[np.float64],
     MSD: npt.NDArray[np.float64],
     beta: npt.NDArray[np.float64],
@@ -151,11 +108,11 @@ def get_gamma(
     Get the gammas using downward recursion from the root nodes.
     The conditional probability of states, given observation of the whole tree P(z_n = k | X_bar = x_bar)
     x_bar is the observations for the whole tree.
-    gamma_1 (k) = P(z_1 = k | X_bar = x_bar)
-    gamma_n (k) = P(z_n = k | X_bar = x_bar)
 
+    :param tree: CSR array representing the lineage tree adjacency
+    :param T: State transitions matrix
     :param MSD: The marginal state distribution P(z_n = k)
-    :param betas: beta values. The conditional probability of states, given observations of the sub-tree rooted in cell_n
+    :param beta: beta values. The conditional probability of states, given observations of the sub-tree rooted in cell_n
     """
     gamma = np.zeros_like(beta)
     gamma[0, :] = beta[0, :]
@@ -165,14 +122,11 @@ def get_gamma(
     coeffs = np.maximum(coeffs, epss)
     beta_parents = T @ coeffs.T
 
-    # Getting lineage by generation, but it is sorted this way
-    for pidx, cis in enumerate(cell_to_daughters):
-        for ci in cis:
-            if ci == -1:
-                continue
-
-            A = gamma[pidx, :].T / beta_parents[:, ci]
-
+    non_leaves = np.nonzero(np.diff(tree.indptr) > 0)[0]
+    for pidx in non_leaves:
+        A_base = gamma[pidx, :].T
+        for ci in tree.indices[tree.indptr[pidx] : tree.indptr[pidx + 1]]:
+            A = A_base / beta_parents[:, ci]
             gamma[ci, :] = coeffs[ci, :] * (A @ T)
 
     assert np.all(np.isfinite(gamma))