phylogenie 2.0.14__tar.gz → 2.1.1__tar.gz

{phylogenie-2.0.14 → phylogenie-2.1.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: phylogenie
-Version: 2.0.14
+Version: 2.1.1
 Summary: Generate phylogenetic datasets with minimal setup effort
 Author: Gabriele Marino
 Author-email: gabmarino.8601@gmail.com
@@ -24,7 +24,6 @@ Description-Content-Type: text/markdown
 
 [![AliSim](https://img.shields.io/badge/Powered%20by-AliSim-orange?style=flat-square)](https://iqtree.github.io/doc/AliSim)
 [![PyPI version](https://img.shields.io/pypi/v/phylogenie)](https://pypi.org/project/phylogenie/)
-[![PyPI downloads](https://shields.io/pypi/dm/phylogenie)](https://pypi.org/project/phylogenie/)
 
 Phylogenie is a [Python](https://www.python.org/) package designed to easily simulate phylogenetic datasets—such as trees and multiple sequence alignments (MSAs)—with minimal setup effort. Simply specify the distributions from which your parameters should be sampled, and Phylogenie will handle the rest!
 

{phylogenie-2.0.14 → phylogenie-2.1.1}/README.md

@@ -6,7 +6,6 @@
 
 [![AliSim](https://img.shields.io/badge/Powered%20by-AliSim-orange?style=flat-square)](https://iqtree.github.io/doc/AliSim)
 [![PyPI version](https://img.shields.io/pypi/v/phylogenie)](https://pypi.org/project/phylogenie/)
-[![PyPI downloads](https://shields.io/pypi/dm/phylogenie)](https://pypi.org/project/phylogenie/)
 
 Phylogenie is a [Python](https://www.python.org/) package designed to easily simulate phylogenetic datasets—such as trees and multiple sequence alignments (MSAs)—with minimal setup effort. Simply specify the distributions from which your parameters should be sampled, and Phylogenie will handle the rest!
 

{phylogenie-2.0.14 → phylogenie-2.1.1}/phylogenie/__init__.py

@@ -8,12 +8,9 @@ from phylogenie.generators import (
     DatasetGeneratorConfig,
     EpidemiologicalTreeDatasetGenerator,
     FBDTreeDatasetGenerator,
-    SkylineMatrixModel,
-    SkylineParameterModel,
-    SkylineVectorModel,
     TreeDatasetGeneratorConfig,
 )
-from phylogenie.io import load_fasta, load_newick
+from phylogenie.io import dump_newick, load_fasta, load_newick
 from phylogenie.msa import MSA
 from phylogenie.skyline import (
     SkylineMatrix,
@@ -29,11 +26,21 @@ from phylogenie.skyline import (
 )
 from phylogenie.tree import Tree
 from phylogenie.treesimulator import (
+    Birth,
+    BirthWithContactTracing,
+    Death,
     Event,
+    Migration,
+    Mutation,
+    MutationTargetType,
+    Sampling,
+    SamplingWithContactTracing,
+    generate_trees,
     get_BD_events,
     get_BDEI_events,
     get_BDSS_events,
     get_canonical_events,
+    get_contact_tracing_events,
     get_epidemiological_events,
     get_FBD_events,
     simulate_tree,
@@ -51,27 +58,35 @@ __all__ = [
     "FBDTreeDatasetGenerator",
     "SkylineMatrix",
     "SkylineMatrixCoercible",
-    "SkylineMatrixModel",
     "skyline_matrix",
     "SkylineParameter",
     "SkylineParameterLike",
-    "SkylineParameterModel",
     "skyline_parameter",
     "SkylineVector",
     "SkylineVectorCoercible",
     "SkylineVectorLike",
-    "SkylineVectorModel",
     "skyline_vector",
     "Tree",
     "TreeDatasetGeneratorConfig",
+    "Birth",
+    "BirthWithContactTracing",
+    "Death",
     "Event",
+    "Migration",
+    "Mutation",
+    "MutationTargetType",
+    "Sampling",
+    "SamplingWithContactTracing",
     "get_BD_events",
     "get_BDEI_events",
     "get_BDSS_events",
     "get_canonical_events",
+    "get_contact_tracing_events",
     "get_epidemiological_events",
     "get_FBD_events",
+    "generate_trees",
     "simulate_tree",
+    "dump_newick",
     "load_fasta",
     "load_newick",
     "MSA",

{phylogenie-2.0.14 → phylogenie-2.1.1}/phylogenie/generators/__init__.py

@@ -3,11 +3,6 @@ from typing import Annotated
 from pydantic import Field
 
 from phylogenie.generators.alisim import AliSimDatasetGenerator
-from phylogenie.generators.configs import (
-    SkylineMatrixModel,
-    SkylineParameterModel,
-    SkylineVectorModel,
-)
 from phylogenie.generators.dataset import DatasetGenerator
 from phylogenie.generators.trees import (
     BDEITreeDatasetGenerator,
@@ -34,7 +29,4 @@ __all__ = [
     "BDTreeDatasetGenerator",
     "BDEITreeDatasetGenerator",
     "BDSSTreeDatasetGenerator",
-    "SkylineMatrixModel",
-    "SkylineParameterModel",
-    "SkylineVectorModel",
 ]

{phylogenie-2.0.14 → phylogenie-2.1.1}/phylogenie/generators/configs.py

@@ -1,12 +1,6 @@
-from pydantic import BaseModel, ConfigDict
-
 import phylogenie.typings as pgt
-
-
-class Distribution(BaseModel):
-    type: str
-    model_config = ConfigDict(extra="allow")
-
+from phylogenie.treesimulator import MutationTargetType
+from phylogenie.utils import Distribution, StrictBaseModel
 
 Integer = str | int
 Scalar = str | pgt.Scalar
@@ -15,10 +9,6 @@ OneOrManyScalars = Scalar | pgt.Many[Scalar]
 OneOrMany2DScalars = Scalar | pgt.Many2D[Scalar]
 
 
-class StrictBaseModel(BaseModel):
-    model_config = ConfigDict(extra="forbid")
-
-
 class SkylineParameterModel(StrictBaseModel):
     value: ManyScalars
     change_times: ManyScalars
@@ -37,3 +27,12 @@ class SkylineMatrixModel(StrictBaseModel):
 SkylineParameter = Scalar | SkylineParameterModel
 SkylineVector = str | pgt.Scalar | pgt.Many[SkylineParameter] | SkylineVectorModel
 SkylineMatrix = str | pgt.Scalar | pgt.Many[SkylineVector] | SkylineMatrixModel | None
+
+
+class Event(StrictBaseModel):
+    states: str | list[str] | None = None
+    rate: SkylineParameter
+
+
+class Mutation(Event):
+    rate_scalers: dict[MutationTargetType, Distribution]

{phylogenie-2.0.14 → phylogenie-2.1.1}/phylogenie/generators/dataset.py

@@ -10,7 +10,8 @@ import pandas as pd
 from numpy.random import Generator, default_rng
 from tqdm import tqdm
 
-import phylogenie.generators.configs as cfg
+from phylogenie.generators.factories import distribution
+from phylogenie.utils import Distribution, StrictBaseModel
 
 
 class DataType(str, Enum):
@@ -22,12 +23,12 @@ DATA_DIRNAME = "data"
 METADATA_FILENAME = "metadata.csv"
 
 
-class DatasetGenerator(ABC, cfg.StrictBaseModel):
+class DatasetGenerator(ABC, StrictBaseModel):
     output_dir: str = "phylogenie-outputs"
     n_samples: int | dict[str, int] = 1
     n_jobs: int = -1
     seed: int | None = None
-    context: dict[str, cfg.Distribution] | None = None
+    context: dict[str, Distribution] | None = None
 
     @abstractmethod
     def _generate_one(
@@ -55,19 +56,21 @@ class DatasetGenerator(ABC, cfg.StrictBaseModel):
         data: list[dict[str, Any]] = [{} for _ in range(n_samples)]
         if self.context is not None:
             for d, (k, v) in product(data, self.context.items()):
-                args = v.model_extra if v.model_extra is not None else {}
-                d[k] = np.array(getattr(rng, v.type)(**args)).tolist()
+                dist = distribution(v, d)
+                d[k] = np.array(getattr(rng, dist.type)(**dist.args)).tolist()
             df = pd.DataFrame([{"file_id": str(i), **d} for i, d in enumerate(data)])
             df.to_csv(os.path.join(output_dir, METADATA_FILENAME), index=False)
 
-        joblib.Parallel(n_jobs=self.n_jobs)(
+        jobs = joblib.Parallel(n_jobs=self.n_jobs, return_as="generator_unordered")(
             joblib.delayed(self.generate_one)(
                 filename=os.path.join(data_dir, str(i)),
                 data=data[i],
                 seed=int(rng.integers(2**32)),
             )
-            for i in tqdm(range(n_samples), desc=f"Generating {data_dir}...")
+            for i in range(n_samples)
         )
+        for _ in tqdm(jobs, total=n_samples, desc=f"Generating {data_dir}..."):
+            pass
 
     def generate(self) -> None:
         rng = default_rng(self.seed)

{phylogenie-2.0.14 → phylogenie-2.1.1}/phylogenie/generators/factories.py

@@ -14,6 +14,7 @@ from phylogenie.skyline import (
     SkylineVector,
     SkylineVectorCoercible,
 )
+from phylogenie.utils import Distribution
 
 
 def _eval_expression(expression: str, data: dict[str, Any]) -> Any:
@@ -29,6 +30,14 @@ def _eval_expression(expression: str, data: dict[str, Any]) -> Any:
     ).tolist()
 
 
+def distribution(x: Distribution, data: dict[str, Any]) -> Distribution:
+    args = x.args
+    for arg_name, arg_value in args.items():
+        if isinstance(arg_value, str):
+            args[arg_name] = _eval_expression(arg_value, data)
+    return Distribution(type=x.type, **args)
+
+
 def integer(x: cfg.Integer, data: dict[str, Any]) -> int:
     if isinstance(x, str):
         e = _eval_expression(x, data)

{phylogenie-2.0.14 → phylogenie-2.1.1}/phylogenie/generators/trees.py

@@ -9,6 +9,7 @@ from pydantic import Field
 import phylogenie.generators.configs as cfg
 from phylogenie.generators.dataset import DatasetGenerator, DataType
 from phylogenie.generators.factories import (
+    distribution,
     integer,
     scalar,
     skyline_matrix,
@@ -19,10 +20,12 @@ from phylogenie.io import dump_newick
 from phylogenie.tree import Tree
 from phylogenie.treesimulator import (
     Event,
+    Mutation,
     get_BD_events,
     get_BDEI_events,
     get_BDSS_events,
     get_canonical_events,
+    get_contact_tracing_events,
     get_epidemiological_events,
     get_FBD_events,
     simulate_tree,
@@ -41,16 +44,13 @@ class ParameterizationType(str, Enum):
 class TreeDatasetGenerator(DatasetGenerator):
     data_type: Literal[DataType.TREES] = DataType.TREES
     min_tips: cfg.Integer = 1
-    max_tips: cfg.Integer | None = None
+    max_tips: cfg.Integer = 2**32
     max_time: cfg.Scalar = np.inf
     init_state: str | None = None
     sampling_probability_at_present: cfg.Scalar = 0.0
-    max_tries: int | None = None
-    notification_probability: cfg.Scalar = 0.0
-    max_notified_contacts: cfg.Integer = 1
-    samplable_states_after_notification: list[str] | None = None
-    sampling_rate_after_notification: cfg.SkylineParameter = np.inf
-    contacts_removal_probability: cfg.SkylineParameter = 1
+
+    @abstractmethod
+    def _get_events(self, data: dict[str, Any]) -> list[Event]: ...
 
     def simulate_one(self, rng: Generator, data: dict[str, Any]) -> Tree | None:
         events = self._get_events(data)
@@ -59,34 +59,18 @@ class TreeDatasetGenerator(DatasetGenerator):
             if self.init_state is None
             else self.init_state.format(**data)
         )
-        max_tips = (
-            self.max_tips if self.max_tips is None else integer(self.max_tips, data)
-        )
         return simulate_tree(
             events=events,
             min_tips=integer(self.min_tips, data),
-            max_tips=max_tips,
+            max_tips=integer(self.max_tips, data),
             max_time=scalar(self.max_time, data),
             init_state=init_state,
             sampling_probability_at_present=scalar(
                 self.sampling_probability_at_present, data
             ),
-            notification_probability=scalar(self.notification_probability, data),
-            max_notified_contacts=integer(self.max_notified_contacts, data),
-            samplable_states_after_notification=self.samplable_states_after_notification,
-            sampling_rate_after_notification=skyline_parameter(
-                self.sampling_rate_after_notification, data
-            ),
-            contacts_removal_probability=skyline_parameter(
-                self.contacts_removal_probability, data
-            ),
-            max_tries=self.max_tries,
             seed=int(rng.integers(2**32)),
         )
 
-    @abstractmethod
-    def _get_events(self, data: dict[str, Any]) -> list[Event]: ...
-
     def _generate_one(
         self, filename: str, rng: Generator, data: dict[str, Any]
     ) -> None:
@@ -101,9 +85,9 @@ class CanonicalTreeDatasetGenerator(TreeDatasetGenerator):
     )
     states: list[str]
     sampling_rates: cfg.SkylineVector
+    remove_after_sampling: bool
     birth_rates: cfg.SkylineVector = 0
     death_rates: cfg.SkylineVector = 0
-    removal_probabilities: cfg.SkylineVector = 0
     migration_rates: cfg.SkylineMatrix = None
     birth_rates_among_states: cfg.SkylineMatrix = None
 
@@ -111,9 +95,9 @@ class CanonicalTreeDatasetGenerator(TreeDatasetGenerator):
         return get_canonical_events(
             states=self.states,
             sampling_rates=skyline_vector(self.sampling_rates, data),
+            remove_after_sampling=self.remove_after_sampling,
             birth_rates=skyline_vector(self.birth_rates, data),
             death_rates=skyline_vector(self.death_rates, data),
-            removal_probabilities=skyline_vector(self.removal_probabilities, data),
             migration_rates=skyline_matrix(self.migration_rates, data),
             birth_rates_among_states=skyline_matrix(
                 self.birth_rates_among_states, data
@@ -121,19 +105,87 @@ class CanonicalTreeDatasetGenerator(TreeDatasetGenerator):
         )
 
 
-class EpidemiologicalTreeDatasetGenerator(TreeDatasetGenerator):
+class FBDTreeDatasetGenerator(TreeDatasetGenerator):
+    parameterization: Literal[ParameterizationType.FBD] = ParameterizationType.FBD
+    states: list[str]
+    sampling_proportions: cfg.SkylineVector
+    diversification: cfg.SkylineVector = 0
+    turnover: cfg.SkylineVector = 0
+    migration_rates: cfg.SkylineMatrix = None
+    diversification_between_states: cfg.SkylineMatrix = None
+
+    def _get_events(self, data: dict[str, Any]) -> list[Event]:
+        return get_FBD_events(
+            states=self.states,
+            diversification=skyline_vector(self.diversification, data),
+            turnover=skyline_vector(self.turnover, data),
+            sampling_proportions=skyline_vector(self.sampling_proportions, data),
+            migration_rates=skyline_matrix(self.migration_rates, data),
+            diversification_between_states=skyline_matrix(
+                self.diversification_between_states, data
+            ),
+        )
+
+
+class TreeDatasetGeneratorForEpidemiology(TreeDatasetGenerator):
+    max_notified_contacts: cfg.Integer = 1
+    notification_probability: cfg.SkylineParameter = 0.0
+    sampling_rate_after_notification: cfg.SkylineParameter = np.inf
+    samplable_states_after_notification: list[str] | None = None
+    mutations: tuple[cfg.Mutation, ...] = Field(default_factory=tuple)
+
+    @abstractmethod
+    def _get_base_events(self, data: dict[str, Any]) -> list[Event]: ...
+
+    def _get_events(self, data: dict[str, Any]) -> list[Event]:
+        events = self._get_base_events(data)
+        if self.notification_probability:
+            events = get_contact_tracing_events(
+                events=events,
+                max_notified_contacts=integer(self.max_notified_contacts, data),
+                notification_probability=skyline_parameter(
+                    self.notification_probability, data
+                ),
+                sampling_rate_after_notification=skyline_parameter(
+                    self.sampling_rate_after_notification, data
+                ),
+                samplable_states_after_notification=self.samplable_states_after_notification,
+            )
+        all_states = list({e.state for e in events})
+        for mutation in self.mutations:
+            states = mutation.states
+            if isinstance(states, str):
+                states = [states]
+            elif states is None:
+                states = all_states
+            for state in states:
+                if state not in all_states:
+                    raise ValueError(
+                        f"Mutation state '{state}' is not found in states {all_states}."
+                    )
+                rate_scalers = {
+                    t: distribution(r, data) for t, r in mutation.rate_scalers.items()
+                }
+                events.append(
+                    Mutation(
+                        state, skyline_parameter(mutation.rate, data), rate_scalers
+                    )
+                )
+        return events
+
+
+class EpidemiologicalTreeDatasetGenerator(TreeDatasetGeneratorForEpidemiology):
     parameterization: Literal[ParameterizationType.EPIDEMIOLOGICAL] = (
         ParameterizationType.EPIDEMIOLOGICAL
     )
     states: list[str]
+    sampling_proportions: cfg.SkylineVector
     reproduction_numbers: cfg.SkylineVector = 0
     become_uninfectious_rates: cfg.SkylineVector = 0
-    sampling_proportions: cfg.SkylineVector = 1
-    removal_probabilities: cfg.SkylineVector = 1
     migration_rates: cfg.SkylineMatrix = None
     reproduction_numbers_among_states: cfg.SkylineMatrix = None
 
-    def _get_events(self, data: dict[str, Any]) -> list[Event]:
+    def _get_base_events(self, data: dict[str, Any]) -> list[Event]:
         return get_epidemiological_events(
             states=self.states,
             reproduction_numbers=skyline_vector(self.reproduction_numbers, data),
@@ -141,7 +193,6 @@ class EpidemiologicalTreeDatasetGenerator(TreeDatasetGenerator):
                 self.become_uninfectious_rates, data
             ),
             sampling_proportions=skyline_vector(self.sampling_proportions, data),
-            removal_probabilities=skyline_vector(self.removal_probabilities, data),
             migration_rates=skyline_matrix(self.migration_rates, data),
             reproduction_numbers_among_states=skyline_matrix(
                 self.reproduction_numbers_among_states, data
@@ -149,37 +200,13 @@ class EpidemiologicalTreeDatasetGenerator(TreeDatasetGenerator):
         )
 
 
-class FBDTreeDatasetGenerator(TreeDatasetGenerator):
-    parameterization: Literal[ParameterizationType.FBD] = ParameterizationType.FBD
-    states: list[str]
-    diversification: cfg.SkylineVector = 0
-    turnover: cfg.SkylineVector = 0
-    sampling_proportions: cfg.SkylineVector = 1
-    removal_probabilities: cfg.SkylineVector = 0
-    migration_rates: cfg.SkylineMatrix = None
-    diversification_between_types: cfg.SkylineMatrix = None
-
-    def _get_events(self, data: dict[str, Any]) -> list[Event]:
-        return get_FBD_events(
-            states=self.states,
-            diversification=skyline_vector(self.diversification, data),
-            turnover=skyline_vector(self.turnover, data),
-            sampling_proportions=skyline_vector(self.sampling_proportions, data),
-            removal_probabilities=skyline_vector(self.removal_probabilities, data),
-            migration_rates=skyline_matrix(self.migration_rates, data),
-            diversification_between_types=skyline_matrix(
-                self.diversification_between_types, data
-            ),
-        )
-
-
-class BDTreeDatasetGenerator(TreeDatasetGenerator):
+class BDTreeDatasetGenerator(TreeDatasetGeneratorForEpidemiology):
     parameterization: Literal[ParameterizationType.BD] = ParameterizationType.BD
     reproduction_number: cfg.SkylineParameter
     infectious_period: cfg.SkylineParameter
     sampling_proportion: cfg.SkylineParameter = 1
 
-    def _get_events(self, data: dict[str, Any]) -> list[Event]:
+    def _get_base_events(self, data: dict[str, Any]) -> list[Event]:
         return get_BD_events(
             reproduction_number=skyline_parameter(self.reproduction_number, data),
             infectious_period=skyline_parameter(self.infectious_period, data),
@@ -187,14 +214,14 @@ class BDTreeDatasetGenerator(TreeDatasetGenerator):
         )
 
 
-class BDEITreeDatasetGenerator(TreeDatasetGenerator):
+class BDEITreeDatasetGenerator(TreeDatasetGeneratorForEpidemiology):
     parameterization: Literal[ParameterizationType.BDEI] = ParameterizationType.BDEI
     reproduction_number: cfg.SkylineParameter
     infectious_period: cfg.SkylineParameter
     incubation_period: cfg.SkylineParameter
     sampling_proportion: cfg.SkylineParameter = 1
 
-    def _get_events(self, data: dict[str, Any]) -> list[Event]:
+    def _get_base_events(self, data: dict[str, Any]) -> list[Event]:
         return get_BDEI_events(
             reproduction_number=skyline_parameter(self.reproduction_number, data),
             infectious_period=skyline_parameter(self.infectious_period, data),
@@ -203,7 +230,7 @@ class BDEITreeDatasetGenerator(TreeDatasetGenerator):
         )
 
 
-class BDSSTreeDatasetGenerator(TreeDatasetGenerator):
+class BDSSTreeDatasetGenerator(TreeDatasetGeneratorForEpidemiology):
     parameterization: Literal[ParameterizationType.BDSS] = ParameterizationType.BDSS
     reproduction_number: cfg.SkylineParameter
     infectious_period: cfg.SkylineParameter
@@ -211,7 +238,7 @@ class BDSSTreeDatasetGenerator(TreeDatasetGenerator):
     superspreaders_proportion: cfg.SkylineParameter
     sampling_proportion: cfg.SkylineParameter = 1
 
-    def _get_events(self, data: dict[str, Any]) -> list[Event]:
+    def _get_base_events(self, data: dict[str, Any]) -> list[Event]:
         return get_BDSS_events(
             reproduction_number=skyline_parameter(self.reproduction_number, data),
             infectious_period=skyline_parameter(self.infectious_period, data),

{phylogenie-2.0.14 → phylogenie-2.1.1}/phylogenie/io.py

@@ -24,21 +24,28 @@ def _parse_newick(newick: str) -> Tree:
             stack.append(current_nodes)
             current_nodes = []
         else:
-            id = _parse_chars([":", ",", ")", ";"])
+            id = _parse_chars([":", ",", ")", ";", "["])
             branch_length = None
             if newick[i] == ":":
                 i += 1
-                branch_length = _parse_chars([",", ")", ";"])
+                branch_length = float(_parse_chars([",", ")", ";", "["]))
 
-            current_node = Tree(
-                id=id,
-                branch_length=(None if branch_length is None else float(branch_length)),
-            )
+            current_node = Tree(id, branch_length)
             for node in current_children:
                 current_node.add_child(node)
                 current_children = []
             current_nodes.append(current_node)
 
+            if newick[i] == "[":
+                i += 1
+                features = _parse_chars(["]"]).split(":")
+                i += 1
+                if features[0] != "&&NHX":
+                    raise ValueError(f"Expected '&&NHX' for node features.")
+                for feature in features[1:]:
+                    key, value = feature.split("=", 1)
+                    current_node.set(key, eval(value))
+
             if newick[i] == ")":
                 current_children = current_nodes
                 current_nodes = stack.pop()
@@ -47,7 +54,7 @@ def _parse_newick(newick: str) -> Tree:
 
         i += 1
 
-    raise ValueError("Newick string does not end with a semicolon.")
+    raise ValueError("Newick string is invalid.")
 
 
 def load_newick(filepath: str) -> Tree | list[Tree]:
@@ -63,12 +70,19 @@ def _to_newick(tree: Tree) -> str:
         newick = f"({children_newick}){newick}"
     if tree.branch_length is not None:
         newick += f":{tree.branch_length}"
+    if tree.features:
+        reprs = {k: repr(v).replace("'", '"') for k, v in tree.features.items()}
+        features = [f"{k}={repr}" for k, repr in reprs.items()]
+        newick += f"[&&NHX:{':'.join(features)}]"
     return newick
 
 
-def dump_newick(tree: Tree, filepath: str) -> None:
+def dump_newick(trees: Tree | list[Tree], filepath: str) -> None:
+    if isinstance(trees, Tree):
+        trees = [trees]
     with open(filepath, "w") as file:
-        file.write(_to_newick(tree) + ";")
+        for t in trees:
+            file.write(_to_newick(t) + ";\n")
 
 
 def load_fasta(
@@ -80,13 +94,14 @@ def load_fasta(
             if not line.startswith(">"):
                 raise ValueError(f"Invalid FASTA format: expected '>', got '{line[0]}'")
             id = line[1:].strip()
+            time = None
             if extract_time_from_id is not None:
                 time = extract_time_from_id(id)
-            else:
+            elif "|" in id:
                 try:
                     time = float(id.split("|")[-1])
-                except:
-                    time = None
+                except ValueError:
+                    pass
             chars = next(f).strip()
             sequences.append(Sequence(id, chars, time))
     return MSA(sequences)

{phylogenie-2.0.14 → phylogenie-2.1.1}/phylogenie/tree.py

@@ -1,18 +1,41 @@
 from collections.abc import Iterator
+from typing import Any
 
 
 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] = []
+        self._parent: Tree | None = None
+        self._children: list[Tree] = []
+        self._features: dict[str, Any] = {}
+
+    @property
+    def children(self) -> tuple["Tree", ...]:
+        return tuple(self._children)
+
+    @property
+    def parent(self) -> "Tree | None":
+        return self._parent
+
+    @property
+    def features(self) -> dict[str, Any]:
+        return self._features
 
     def add_child(self, child: "Tree") -> "Tree":
-        child.parent = self
-        self.children.append(child)
+        child._parent = self
+        self._children.append(child)
         return self
 
+    def remove_child(self, child: "Tree") -> None:
+        self._children.remove(child)
+        child._parent = None
+
+    def set_parent(self, node: "Tree | None"):
+        self._parent = node
+        if node is not None:
+            node._children.append(self)
+
     def preorder_traversal(self) -> Iterator["Tree"]:
         yield self
         for child in self.children:
@@ -29,6 +52,9 @@ class Tree:
                 return node
         raise ValueError(f"Node with id {id} not found.")
 
+    def is_leaf(self) -> bool:
+        return not self.children
+
     def get_leaves(self) -> list["Tree"]:
         return [node for node in self if not node.children]
 
@@ -38,17 +64,17 @@ class Tree:
             raise ValueError(f"Branch length of node {self.id} is not set.")
         return self.branch_length + parent_time
 
-    def is_leaf(self) -> bool:
-        return not self.children
+    def set(self, key: str, value: Any) -> None:
+        self._features[key] = value
 
-    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 get(self, key: str) -> Any:
+        return self._features.get(key)
+
+    def delete(self, key: str) -> None:
+        del self._features[key]
 
     def __iter__(self) -> Iterator["Tree"]:
         return self.preorder_traversal()
 
     def __repr__(self) -> str:
-        return f"TreeNode(id='{self.id}', branch_length={self.branch_length})"
+        return f"TreeNode(id='{self.id}', branch_length={self.branch_length}, features={self.features})"