Likelihood

src/pmhn/_trees/_backend.py

@@ -2,9 +2,11 @@
 
 
 import numpy as np
-
-
+from scipy.sparse.linalg import spsolve_triangular
+from scipy.sparse import csr_matrix
 from pmhn._trees._interfaces import Tree
+from pmhn._trees._tree_utils import create_all_subtrees, bfs_compare
+from anytree import Node
 
 
 class IndividualTreeMHNBackendInterface(Protocol):
@@ -57,26 +59,136 @@ def gradient_and_loglikelihood(
 
 
 class OriginalTreeMHNBackend(IndividualTreeMHNBackendInterface):
+    def __init__(self, jitter: float = 1e-10):
+        self._jitter = jitter
+
+    _jitter: float
+
+    def create_V_Mat(
+        self, tree: Node, theta: np.ndarray, sampling_rate: float
+    ) -> np.ndarray:
+        """Calculates the V matrix.
+
+        Args:
+            tree: a tree
+            theta: real-valued (i.e., log-theta) matrix,
+              shape (n_mutations, n_mutations)
+            sampling_rate: a scalar of type float
+        Returns:
+            the V matrix.
+        """
+
+        subtrees = create_all_subtrees(tree)
+        subtrees_size = len(subtrees)
+        Q = np.zeros((subtrees_size, subtrees_size))
+        for i in range(subtrees_size):
+            for j in range(subtrees_size):
+                if i == j:
+                    Q[i][j] = self.diag_entry(subtrees[i], theta)
+                else:
+                    Q[i][j] = self.off_diag_entry(subtrees[i], subtrees[j], theta)
+        V = np.eye(subtrees_size) * sampling_rate - Q
+        return V
+
+    def diag_entry(self, tree: Node, theta: np.ndarray) -> float:
+        """
+        Calculates a diagonal entry of the V matrix.
+
+        Args:
+            tree: a tree
+            theta: real-valued (i.e., log-theta) matrix,
+              shape (n_mutations, n_mutations)
+
+        Returns:
+            the diagonal entry of the V matrix corresponding to tree
+        """
+        lamb_sum = 0
+        n_mutations = len(theta)
+        current_nodes = [tree]
+        while len(current_nodes) != 0:
+            next_nodes = []
+            for node in current_nodes:
+                tree_mutations = list(node.path) + list(node.children)
+                exit_mutations = list(
+                    set([i + 1 for i in range(n_mutations)]).difference(
+                        set(
+                            [
+                                tree_mutation.name  # type: ignore
+                                for tree_mutation in tree_mutations
+                            ]
+                        )
+                    )
+                )
+                for mutation in exit_mutations:
+                    lamb = 0
+                    exit_subclone = [
+                        anc.name  # type: ignore
+                        for anc in node.path
+                        if anc.parent is not None
+                    ] + [mutation]
+                    for j in exit_subclone:
+                        lamb += theta[mutation - 1][j - 1]
+                    lamb = np.exp(lamb)
+                    lamb_sum -= lamb
+                for child in node.children:
+                    next_nodes.append(child)
+            current_nodes = next_nodes
+        return lamb_sum
+
+    def off_diag_entry(self, tree1: Node, tree2: Node, theta: np.ndarray) -> float:
+        """
+        Calculates an off-diagonal entry of the V matrix.
+
+        Args:
+            tree1: the first tree
+            tree2: the second tree
+            theta: real-valued (i.e., log-theta) matrix,
+              shape (n_mutations, n_mutations)
+
+        Returns:
+            the off-diagonal entry of the V matrix corresponding to tree1 and tree2
+        """
+        exit_node = bfs_compare(tree1, tree2)
+        lamb = 0
+        if exit_node is None:
+            return lamb
+        else:
+            for j in [
+                node.name  # type: ignore
+                for node in exit_node.path
+                if node.parent is not None
+            ]:
+                lamb += theta[exit_node.name - 1][j - 1]
+            lamb = np.exp(lamb)
+            return float(lamb)
+
     def loglikelihood(
-        self,
-        tree: Tree,
-        theta: np.ndarray,
+        self, tree: Node, theta: np.ndarray, sampling_rate: float
     ) -> float:
-        """Calculates loglikelihood `log P(tree | theta)`.
+        """
+        Calculates loglikelihood `log P(tree | theta)`.
 
         Args:
             tree: a tree
             theta: real-valued (i.e., log-theta) matrix,
               shape (n_mutations, n_mutations)
-
+            sampling_rate: a scalar of type float
         Returns:
             loglikelihood of the tree
         """
         # TODO(Pawel): this is part of https://github.com/cbg-ethz/pMHN/issues/15
         #   It can be implemented in any way.
-        raise NotImplementedError
+        V = self.create_V_Mat(tree=tree, theta=theta, sampling_rate=sampling_rate)
+        V_size = V.shape[0]
+        b = np.zeros(V_size)
+        b[0] = 1
+        V_transposed = V.transpose()
+        V_csr = csr_matrix(V_transposed)
+        x = spsolve_triangular(V_csr, b, lower=True)
+
+        return np.log(x[V_size - 1] + self._jitter) + np.log(sampling_rate)
 
-    def gradient(self, tree: Tree, theta: np.ndarray) -> np.ndarray:
+    def gradient(self, tree: Node, theta: np.ndarray) -> np.ndarray:
         """Calculates the partial derivatives of `log P(tree | theta)`
         with respect to `theta`.
 

src/pmhn/_trees/_io.py

@@ -33,6 +33,7 @@ class TreeNaming:
 
     node: str = "Node_ID"
     parent: str = "Parent_ID"
+    mutation: str = "Mutation_ID"
     data: dict[str, str] = dataclasses.field(
         default_factory=lambda: {"Mutation_ID": "mutation"}
     )
@@ -52,7 +53,7 @@ class ForestNaming:
 
 
 def parse_tree(df: pd.DataFrame, naming: TreeNaming) -> anytree.Node:
-    """Parses a data frame into a tree.
+    """Parses a data frame into a tree
 
     Args:
         df: data frame with columns specified in `naming`.
@@ -79,10 +80,12 @@ def parse_tree(df: pd.DataFrame, naming: TreeNaming) -> anytree.Node:
                     f"Root is {root}, but {node_id} == {parent_id} "
                     "also looks like a root."
                 )
-            root = anytree.Node(node_id, parent=None, **values)
+            root = anytree.Node(row[naming.mutation], parent=None, **values)
             nodes[node_id] = root
         else:
-            nodes[node_id] = anytree.Node(node_id, parent=nodes[parent_id], **values)
+            nodes[node_id] = anytree.Node(
+                row[naming.mutation], parent=nodes[parent_id], **values
+            )
 
     if root is None:
         raise ValueError("No root found.")

src/pmhn/_trees/_simulate.py

@@ -67,7 +67,7 @@ def _find_possible_mutations(old_mutations: list[int], n_mutations: int) -> list
             )
 
     possible_mutations = list(
-        set([i + 1 for i in range(n_mutations)]).difference(set(old_mutations))
+        set(range(1, n_mutations + 1)).difference(set(old_mutations))
     )
     return possible_mutations