phylogenie 1.0.8__py3-none-any.whl → 2.0.1__py3-none-any.whl

phylogenie/io.py

@@ -0,0 +1,92 @@
+from typing import Callable
+
+from phylogenie.msa import MSA, Sequence
+from phylogenie.tree import Tree
+
+
+def _parse_newick(newick: str) -> Tree:
+    newick = newick.strip()
+    stack: list[list[Tree]] = []
+    current_children: list[Tree] = []
+    current_nodes: list[Tree] = []
+    i = 0
+    while i < len(newick):
+
+        def _parse_chars(stoppers: list[str]) -> str:
+            nonlocal i
+            chars = ""
+            while newick[i] not in stoppers:
+                chars += newick[i]
+                i += 1
+            return chars
+
+        if newick[i] == "(":
+            stack.append(current_nodes)
+            current_nodes = []
+        else:
+            id = _parse_chars([":", ",", ")", ";"])
+            branch_length = None
+            if newick[i] == ":":
+                i += 1
+                branch_length = _parse_chars([",", ")", ";"])
+
+            current_node = Tree(
+                id=id,
+                branch_length=(None if branch_length is None else float(branch_length)),
+            )
+            for node in current_children:
+                current_node.add_child(node)
+                current_children = []
+            current_nodes.append(current_node)
+
+            if newick[i] == ")":
+                current_children = current_nodes
+                current_nodes = stack.pop()
+            elif newick[i] == ";":
+                return current_node
+
+        i += 1
+
+    raise ValueError("Newick string does not end with a semicolon.")
+
+
+def load_newick(filepath: str) -> Tree | list[Tree]:
+    with open(filepath, "r") as file:
+        trees = [_parse_newick(newick) for newick in file]
+    return trees[0] if len(trees) == 1 else trees
+
+
+def _to_newick(tree: Tree) -> str:
+    children_newick = ",".join([_to_newick(child) for child in tree.children])
+    newick = tree.id
+    if children_newick:
+        newick = f"({children_newick}){newick}"
+    if tree.branch_length is not None:
+        newick += f":{tree.branch_length}"
+    return newick
+
+
+def dump_newick(tree: Tree, filepath: str) -> None:
+    with open(filepath, "w") as file:
+        file.write(_to_newick(tree) + ";")
+
+
+def load_fasta(
+    fasta_file: str, extract_time_from_id: Callable[[str], float] | None = None
+) -> MSA:
+    sequences: list[Sequence] = []
+    with open(fasta_file, "r") as f:
+        for line in f:
+            if not line.startswith(">"):
+                raise ValueError(f"Invalid FASTA format: expected '>', got '{line[0]}'")
+            id = line[1:].strip()
+            if extract_time_from_id is not None:
+                time = extract_time_from_id(id)
+            else:
+                try:
+                    time = float(id.split("|")[-1])
+                except ValueError:
+                    time = None
+            chars = next(f).strip()
+            sequences.append(Sequence(id, chars, time))
+    return MSA(sequences)

phylogenie/main.py

@@ -5,8 +5,8 @@ from glob import glob
 from pydantic import TypeAdapter
 from yaml import safe_load
 
-from phylogenie.core import DatasetGeneratorConfig
-from phylogenie.core.dataset import DatasetGenerator
+from phylogenie.generators import DatasetGeneratorConfig
+from phylogenie.generators.dataset import DatasetGenerator
 
 
 def run(config_path: str) -> None:

phylogenie/msa.py

@@ -0,0 +1,72 @@
+from collections.abc import Iterator
+from dataclasses import dataclass
+
+import numpy as np
+
+
+@dataclass
+class Sequence:
+    id: str
+    chars: str
+    time: float | None = None
+
+
+class MSA:
+    def __init__(self, sequences: list[Sequence]):
+        self.sequences = sequences
+        lengths = {len(sequence.chars) for sequence in sequences}
+        if len(lengths) > 1:
+            raise ValueError(
+                f"All sequences in the alignment must have the same length (got lengths: {lengths})"
+            )
+
+    @property
+    def ids(self) -> list[str]:
+        return [sequence.id for sequence in self.sequences]
+
+    @property
+    def times(self) -> list[float]:
+        times: list[float] = []
+        for sequence in self:
+            if sequence.time is None:
+                raise ValueError(f"Time is not set for sequence {sequence.id}.")
+            times.append(sequence.time)
+        return times
+
+    @property
+    def alignment(self) -> list[list[str]]:
+        return [list(sequence.chars) for sequence in self.sequences]
+
+    @property
+    def n_sequences(self) -> int:
+        return len(self.sequences)
+
+    @property
+    def n_sites(self) -> int:
+        return len(self.alignment[0])
+
+    @property
+    def shape(self) -> tuple[int, int]:
+        return self.n_sequences, self.n_sites
+
+    def __len__(self) -> int:
+        return self.n_sequences
+
+    def __getitem__(self, item: int) -> Sequence:
+        return self.sequences[item]
+
+    def __iter__(self) -> Iterator[Sequence]:
+        return iter(self.sequences)
+
+    def count_informative_sites(self) -> int:
+        n_informative_sites = 0
+        for column in np.array(self.alignment).T:
+            column: np.typing.NDArray[np.str_]
+            _, char_counts = np.unique(column, return_counts=True)
+            is_informative_char = char_counts >= 2
+            if (is_informative_char).sum() >= 2:
+                n_informative_sites += 1
+        return n_informative_sites
+
+    def count_unique_sequences(self) -> int:
+        return len(np.unique(self.alignment, axis=0))

phylogenie/skyline/matrix.py

@@ -11,7 +11,6 @@ from phylogenie.skyline.vector import (
     SkylineVectorOperand,
     is_many_skyline_vectors_coercible,
     is_many_skyline_vectors_like,
-    is_skyline_vector_coercible,
     is_skyline_vector_like,
     is_skyline_vector_operand,
     skyline_vector,
@@ -34,17 +33,25 @@ class SkylineMatrix:
     ):
         if params is not None and value is None and change_times is None:
             if is_many_skyline_vectors_like(params):
-                self.params = [skyline_vector(p, len(params)) for p in params]
+                self.params = [
+                    p if isinstance(p, SkylineVector) else SkylineVector(p)
+                    for p in params
+                ]
             else:
                 raise TypeError(
                     f"It is impossible to create a SkylineMatrix from `params` {params} of type {type(params)}. Please provide a sequence composed of SkylineVectorLike objects (a SkylineVectorLike object can either be a SkylineVector or a sequence of scalars and/or SkylineParameters)."
                 )
+            lengths = {len(p) for p in self.params}
+            if len(lengths) > 1:
+                raise ValueError(
+                    f"All `params` must have the same length to create a SkylineMatrix (got params={params} with lengths {lengths})."
+                )
         elif params is None and value is not None and change_times is not None:
             if tg.is_many_3D_scalars(value):
-                matrix_lengths = {len(matrix) for matrix in value}
-                if any(ml != len(value[0]) for ml in matrix_lengths):
+                lengths = {len(matrix) for matrix in value}
+                if len(lengths) > 1:
                     raise ValueError(
-                        f"All matrices in the `value` of a SkylineMatrix must have the same length (got value={value} with matrix lengths={matrix_lengths})."
+                        f"All matrices in the `value` of a SkylineMatrix must have the same number of rows (got matrices={value} with row lengths {lengths})."
                     )
             else:
                 raise TypeError(
@@ -62,9 +69,17 @@ class SkylineMatrix:
             )
 
     @property
-    def N(self) -> int:
+    def n_rows(self) -> int:
         return len(self.params)
 
+    @property
+    def n_cols(self) -> int:
+        return len(self.params[0])
+
+    @property
+    def shape(self) -> tuple[int, int]:
+        return self.n_rows, self.n_cols
+
     @property
     def change_times(self) -> pgt.Vector1D:
         return sorted(set([t for row in self.params for t in row.change_times]))
@@ -83,12 +98,12 @@ class SkylineMatrix:
             [SkylineVector, SkylineVector | SkylineParameter], SkylineVector
         ],
     ) -> "SkylineMatrix":
-        if is_skyline_vector_operand(other):
-            other = skyline_vector(other, N=self.N)
+        if is_skyline_matrix_operand(other):
+            other = skyline_matrix(other, self.n_rows, self.n_cols)
         elif isinstance(other, SkylineMatrix):
-            if other.N != self.N:
+            if other.shape != self.shape:
                 raise ValueError(
-                    f"Expected a SkylineMatrix with the same size as self (N={self.N}), but got {other} with N={other.N}."
+                    f"It is impossible to operate on SkylineMatrices of different shapes (got self={self.shape} and other={other.shape})."
                 )
         else:
             return NotImplemented
@@ -122,7 +137,7 @@ class SkylineMatrix:
 
     @property
     def T(self) -> "SkylineMatrix":
-        return SkylineMatrix([[v[i] for v in self] for i in range(self.N)])
+        return SkylineMatrix([[v[i] for v in self] for i in range(self.n_cols)])
 
     def __bool__(self) -> bool:
         return any(self.params)
@@ -137,52 +152,57 @@ class SkylineMatrix:
         return iter(self.params)
 
     def __len__(self) -> int:
-        return self.N
+        return self.n_rows
 
     @overload
     def __getitem__(self, item: int) -> SkylineVector: ...
     @overload
-    def __getitem__(self, item: slice) -> list[SkylineVector]: ...
+    def __getitem__(self, item: slice) -> "SkylineMatrix": ...
+    @overload
     def __getitem__(
-        self, item: int | slice
-    ) -> Union[SkylineVector, list[SkylineVector]]:
-        return self.params[item]
+        self, item: tuple[int | slice, int | slice]
+    ) -> Union[SkylineParameter, SkylineVector, "SkylineMatrix"]: ...
+    def __getitem__(
+        self, item: int | slice | tuple[int | slice, int | slice]
+    ) -> Union[SkylineParameter | SkylineVector, "SkylineMatrix"]:
+        if isinstance(item, int):
+            return self.params[item]
+        if isinstance(item, slice):
+            return SkylineMatrix(self.params[item])
+        row_idx, col_idx = item
+        if isinstance(row_idx, int):
+            return self.params[row_idx][col_idx]
+        if isinstance(col_idx, int):
+            return SkylineVector([row[col_idx] for row in self.params[row_idx]])
+        return SkylineMatrix([row[col_idx] for row in self.params[row_idx]])
 
     def __setitem__(self, item: int, value: SkylineVectorLike) -> None:
         if not is_skyline_vector_like(value):
             raise TypeError(
-                f"It is impossible to set item {item} of SkylineMatrix with value {value} of type {type(value)}. Please provide a SkylineVectorLike object (i.e., a SkylineVector or a sequence of scalars and/or SkylineParameters)."
+                f"It is impossible to set item of SkylineMatrix to value {value} of type {type(value)}. Please provide a SkylineVectorLike object (i.e., a SkylineVector or a sequence of scalars and/or SkylineParameters)."
             )
-        self.params[item] = skyline_vector(value, N=self.N)
+        self.params[item] = skyline_vector(value, self.n_cols)
 
 
 def skyline_matrix(
-    x: SkylineMatrixCoercible, N: int, zero_diagonal: bool = False
+    x: SkylineMatrixCoercible, n_rows: int, n_cols: int
 ) -> SkylineMatrix:
-    if N <= 0:
+    if n_rows <= 0 or n_cols <= 0:
         raise ValueError(
-            f"N must be a positive integer for SkylineMatrix construction (got N={N})."
+            f" n_rows and n_cols must be positive integers to create a SkylineMatrix (got n_rows={n_rows} and n_cols={n_cols})."
         )
-    if is_skyline_vector_coercible(x):
-        x = SkylineMatrix([[p] * N for p in skyline_vector(x, N)])
-        if zero_diagonal:
-            for i in range(N):
-                x[i][i] = 0
-        return x
-    elif is_many_skyline_vectors_coercible(x):
-        x = SkylineMatrix(
-            [
-                [
-                    (
-                        0
-                        if i == j and is_skyline_parameter_like(v) and zero_diagonal
-                        else p
-                    )
-                    for j, p in enumerate(skyline_vector(v, N))
-                ]
-                for i, v in enumerate(x)
-            ]
+
+    if is_skyline_parameter_like(x):
+        return SkylineMatrix([[x] * n_cols] * n_rows)
+    if is_skyline_vector_like(x) or is_many_skyline_vectors_coercible(x):
+        if len(x) == n_rows:
+            return SkylineMatrix([skyline_vector(p, n_cols) for p in x])
+        elif len(x) == n_cols:
+            return SkylineMatrix([skyline_vector(p, n_rows) for p in x]).T
+        raise ValueError(
+            f"Expected a SkylineVectorLike of size {n_rows} or {n_cols}, got {x} of size {len(x)}."
         )
+
     if not isinstance(x, SkylineMatrix):
         raise TypeError(
             f"It is impossible to coerce {x} of type {type(x)} into a SkylineMatrix. Please provide either:\n"
@@ -191,12 +211,9 @@ def skyline_matrix(
             "- a sequence of SkylineVectorCoercible objects."
         )
 
-    if x.N != N:
+    if x.shape != (n_rows, n_cols):
         raise ValueError(
-            f"Expected an {N}x{N} SkylineMatrix, got {x} of shape {x.N}x{x.N}."
+            f"Expected an SkylineMatrix of shape ({n_rows}, {n_cols}), got {x} of shape {x.shape}."
         )
 
-    if zero_diagonal and any(x[i][i] for i in range(x.N)):
-        raise ValueError(f"Expected a SkylineMatrix with zero diagonal, but got {x}.")
-
     return x

phylogenie/skyline/vector.py

@@ -56,10 +56,10 @@ class SkylineVector:
                 )
         elif params is None and value is not None and change_times is not None:
             if tg.is_many_2D_scalars(value):
-                vector_lengths = {len(vector) for vector in value}
-                if any(vl != len(value[0]) for vl in vector_lengths):
+                lengths = {len(vector) for vector in value}
+                if len(lengths) > 1:
                     raise ValueError(
-                        f"All rows in the `value` of a SkylineVector must have the same length (got value={value} with vector lengths={vector_lengths})."
+                        f"All rows in the `value` of a SkylineVector must have the same length (got value={value} with lengths {lengths})."
                     )
             else:
                 raise TypeError(
@@ -143,10 +143,12 @@ class SkylineVector:
     @overload
     def __getitem__(self, item: int) -> SkylineParameter: ...
     @overload
-    def __getitem__(self, item: slice) -> list[SkylineParameter]: ...
+    def __getitem__(self, item: slice) -> "SkylineVector": ...
     def __getitem__(
         self, item: int | slice
-    ) -> Union[SkylineParameter, list[SkylineParameter]]:
+    ) -> Union[SkylineParameter, "SkylineVector"]:
+        if isinstance(item, slice):
+            return SkylineVector(self.params[item])
         return self.params[item]
 
     def __setitem__(self, item: int, value: SkylineParameterLike) -> None:
@@ -160,7 +162,7 @@ class SkylineVector:
 def skyline_vector(x: SkylineVectorCoercible, N: int) -> SkylineVector:
     if N <= 0:
         raise ValueError(
-            f"N must be a positive integer for SkylineVector construction (got N={N})."
+            f"N must be a positive integer to create a SkylineVector (got N={N})."
         )
     if is_skyline_parameter_like(x):
         return SkylineVector([skyline_parameter(x)] * N)

phylogenie/tree.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+from typing import Iterator
+
+
+class Tree:
+    def __init__(self, id: str = "", branch_length: float | None = None):
+        self.id = id
+        self.branch_length = branch_length
+        self.parent: Tree | None = None
+        self.children: list[Tree] = []
+
+    def add_child(self, child: Tree) -> Tree:
+        child.parent = self
+        self.children.append(child)
+        return self
+
+    def preorder_traversal(self) -> Iterator[Tree]:
+        yield self
+        for child in self.children:
+            yield from child.preorder_traversal()
+
+    def postorder_traversal(self) -> Iterator[Tree]:
+        for child in self.children:
+            yield from child.postorder_traversal()
+        yield self
+
+    def get_node(self, id: str) -> Tree:
+        for node in self:
+            if node.id == id:
+                return node
+        raise ValueError(f"Node with id {id} not found.")
+
+    def get_leaves(self) -> list[Tree]:
+        return [node for node in self if not node.children]
+
+    def get_time(self) -> float:
+        parent_time = 0 if self.parent is None else self.parent.get_time()
+        if self.branch_length is None:
+            raise ValueError(f"Branch length of node {self.id} is not set.")
+        return self.branch_length + parent_time
+
+    def copy(self) -> Tree:
+        new_tree = Tree(self.id, self.branch_length)
+        for child in self.children:
+            new_tree.add_child(child.copy())
+        return new_tree
+
+    def __iter__(self) -> Iterator[Tree]:
+        return self.preorder_traversal()
+
+    def __repr__(self) -> str:
+        return f"TreeNode(id='{self.id}', branch_length={self.branch_length})"

phylogenie/treesimulator/__init__.py

@@ -0,0 +1,21 @@
+from phylogenie.treesimulator.events import (
+    Event,
+    get_BD_events,
+    get_BDEI_events,
+    get_BDSS_events,
+    get_canonical_events,
+    get_epidemiological_events,
+    get_FBD_events,
+)
+from phylogenie.treesimulator.gillespie import simulate_tree
+
+__all__ = [
+    "Event",
+    "get_BD_events",
+    "get_BDEI_events",
+    "get_BDSS_events",
+    "get_canonical_events",
+    "get_epidemiological_events",
+    "get_FBD_events",
+    "simulate_tree",
+]