phylogenie 2.1.26__py3-none-any.whl → 2.1.27__py3-none-any.whl

phylogenie/generators/factories.py

@@ -17,7 +17,7 @@ from phylogenie.skyline import (
     SkylineVector,
     SkylineVectorCoercible,
 )
-from phylogenie.treesimulator import Mutation
+from phylogenie.treesimulator import EventType, Mutation
 
 
 def _eval_expression(expression: str, data: dict[str, Any]) -> Any:
@@ -223,16 +223,21 @@ def distribution(x: Distribution, data: dict[str, Any]) -> Distribution:
 
 
 def mutations(
-    x: list[cfg.Mutation], data: dict[str, Any], states: set[str]
+    x: list[cfg.Mutation],
+    data: dict[str, Any],
+    states: set[str],
+    rates_to_log: list[EventType] | None,
 ) -> list[Mutation]:
     mutations: list[Mutation] = []
     for m in x:
         rate = skyline_parameter(m.rate, data)
         rate_scalers = {k: distribution(v, data) for k, v in m.rate_scalers.items()}
         if m.state is None:
-            mutations.extend(Mutation(s, rate, rate_scalers) for s in states)
+            mutations.extend(
+                Mutation(s, rate, rate_scalers, rates_to_log) for s in states
+            )
         else:
-            mutations.append(Mutation(m.state, rate, rate_scalers))
+            mutations.append(Mutation(m.state, rate, rate_scalers, rates_to_log))
     return mutations
 
 

phylogenie/generators/trees.py

@@ -23,6 +23,7 @@ from phylogenie.models import Distribution
 from phylogenie.tree import Tree
 from phylogenie.treesimulator import (
     Event,
+    EventType,
     Feature,
     get_BD_events,
     get_BDEI_events,
@@ -48,6 +49,7 @@ class ParameterizationType(str, Enum):
 class TreeDatasetGenerator(DatasetGenerator):
     data_type: Literal[DataType.TREES] = DataType.TREES
     mutations: list[cfg.Mutation] = Field(default_factory=lambda: [])
+    rates_to_log: list[EventType] | None = None
     min_tips: cfg.Integer = 1
     max_tips: cfg.Integer | None = None
     max_time: cfg.Scalar = np.inf
@@ -67,9 +69,11 @@ class TreeDatasetGenerator(DatasetGenerator):
             if self.init_state is None
             else self.init_state.format(**data)
         )
-        states = {e.state for e in self._get_events(data)}
+        events = self._get_events(data)
+        states = {e.state for e in events}
+        events += mutations(self.mutations, data, states, self.rates_to_log)
         return simulate_tree(
-            events=self._get_events(data) + mutations(self.mutations, data, states),
+            events=events,
             min_tips=integer(self.min_tips, data),
             max_tips=None if self.max_tips is None else integer(self.max_tips, data),
             max_time=scalar(self.max_time, data),

phylogenie/treesimulator/events/base.py

@@ -15,8 +15,6 @@ class EventType(str, Enum):
     MIGRATION = "migration"
     SAMPLING = "sampling"
     MUTATION = "mutation"
-    BIRTH_WITH_CT = "birth_with_contact_tracing"
-    SAMPLING_WITH_CT = "sampling_with_contact_tracing"
 
 
 class Event(ABC):

phylogenie/treesimulator/events/contact_tracing.py

@@ -23,7 +23,7 @@ def is_CT_state(state: str) -> bool:
 
 
 class BirthWithContactTracing(Event):
-    type = EventType.BIRTH_WITH_CT
+    type = EventType.BIRTH
 
     def __init__(self, state: str, rate: SkylineParameterLike, child_state: str):
         super().__init__(state, rate)
@@ -42,7 +42,7 @@ class BirthWithContactTracing(Event):
 
 
 class SamplingWithContactTracing(Event):
-    type = EventType.SAMPLING_WITH_CT
+    type = EventType.SAMPLING
 
     def __init__(
         self,

phylogenie/treesimulator/events/mutations.py

@@ -43,9 +43,11 @@ class Mutation(Event):
         state: str,
         rate: SkylineParameterLike,
         rate_scalers: dict[EventType, Distribution],
+        rates_to_log: list[EventType] | None = None,
     ):
         super().__init__(state, rate)
         self.rate_scalers = rate_scalers
+        self.rates_to_log = [] if rates_to_log is None else rates_to_log
 
     def apply(
         self, model: Model, events: list[Event], time: float, rng: Generator
@@ -73,18 +75,24 @@ class Mutation(Event):
 
             if isinstance(event, Birth | BirthWithContactTracing):
                 event.child_state = _get_mutated_state(mutation_id, event.child_state)
+                metadata_key = f"birth_rate_from_{event.state}_to_{event.child_state}"
             elif isinstance(event, Migration):
                 event.target_state = _get_mutated_state(mutation_id, event.target_state)
-            elif not isinstance(
+                metadata_key = (
+                    f"migration_rate_from_{event.state}_to_{event.target_state}"
+                )
+            elif isinstance(
                 event, Mutation | Death | Sampling | SamplingWithContactTracing
             ):
+                metadata_key = f"{event.type}_rate_for_{event.state}"
+            else:
                 raise ValueError(
                     f"Mutation not implemented for event of type {type(event)}."
                 )
 
-            if event.type in rate_scalers:
-                event.rate *= rate_scalers[event.type]
-                metadata[f"{MUTATION_PREFIX}{mutation_id}.{event.type}.rate.value"] = (
+            event.rate *= rate_scalers.get(event.type, 1)
+            if event.type in self.rates_to_log:
+                metadata[metadata_key] = (
                     event.rate.value[0]
                     if len(event.rate.value) == 1
                     else list(event.rate.value)

{phylogenie-2.1.26.dist-info → phylogenie-2.1.27.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: phylogenie
-Version: 2.1.26
+Version: 2.1.27
 Summary: Generate phylogenetic datasets with minimal setup effort
 Author: Gabriele Marino
 Author-email: gabmarino.8601@gmail.com

{phylogenie-2.1.26.dist-info → phylogenie-2.1.27.dist-info}/RECORD

@@ -4,8 +4,8 @@ phylogenie/generators/__init__.py,sha256=zsOxy28-9j9alOQLIgrOAFfmM58NNHO_NEtW-KX
 phylogenie/generators/alisim.py,sha256=1YQLpOG_Bpn9YqExQqEu-wz1MDGCbpPmTdhrBb6TbWc,2820
 phylogenie/generators/configs.py,sha256=zjNnljBBC-d6DTjWjieE8-HnKeE58LAy3moBYiD6NFo,1042
 phylogenie/generators/dataset.py,sha256=kY92diePr2IjiLejHLixJoYRc-2LpM-GBt3wkX9SYvA,2109
-phylogenie/generators/factories.py,sha256=6VRYiLoscA1qcTcEVKwsUIitMg3cN-t2C4acxOlab24,8424
-phylogenie/generators/trees.py,sha256=5AHZa3kXizuO583C3LljJloVyzqGd0u5Nz24ZsrhwDE,10062
+phylogenie/generators/factories.py,sha256=qSz-s7ltQwUf7R6TPZ-5xFyOnnSHkKn0u5IukeysXVU,8544
+phylogenie/generators/trees.py,sha256=zsmL3_YUUwtn-_NrGSvwZVSGIm9kpf8SPHxoQCIT_hE,10168
 phylogenie/generators/typeguards.py,sha256=yj4VkhOaUXJ2OrY-6zhOeY9C4yKIQxjZtk2d-vIxttQ,828
 phylogenie/io/__init__.py,sha256=gtRYtDdZSTlWCj3I4vmMJSAs93jdz5RySkCakD3sxlQ,214
 phylogenie/io/fasta.py,sha256=IWtNb_RQLR6kvS0G826wB9SodkCGfugddoUHx78Yrec,837
@@ -22,18 +22,18 @@ phylogenie/skyline/vector.py,sha256=60jtp7PieiEaEH0Tp6zNjNKjyzpN_nT5uwBUXbgeATk,
 phylogenie/tree.py,sha256=P1uM6s32TsODpvNJQIPMix9oj39vGSw_wsHYp2wmy5U,5246
 phylogenie/treesimulator/__init__.py,sha256=0wU1PECAcM5rBg5MmeA3poWgfGDzqHySmSJjJs1wkPw,1088
 phylogenie/treesimulator/events/__init__.py,sha256=w2tJ0D2WB5AiCbr3CsKN6vdADueiAEMzd_ve0rpa4zg,939
-phylogenie/treesimulator/events/base.py,sha256=jklKc6qU3LFhmL84Dxmew0tflRdPdbHYhd_LsU-8QuU,1363
-phylogenie/treesimulator/events/contact_tracing.py,sha256=D0gXKKAIqOjD3z1wZvoGogMlDawpz45AkwC2i7GmgmM,5212
+phylogenie/treesimulator/events/base.py,sha256=Rb9qO0EURKGbFPv_KuRrzZv5Krl1ZEMj8ASHXlBLhbI,1259
+phylogenie/treesimulator/events/contact_tracing.py,sha256=QqinLfQmBBi6qF3nFE_Qsiyd6RUG0RJqxiBOKK0nRGE,5196
 phylogenie/treesimulator/events/core.py,sha256=bhgQgi5L-oaHsoWJmUOsTTzWxi0POYxVLoF-KrC8AGQ,8179
-phylogenie/treesimulator/events/mutations.py,sha256=_SqG_hEwqiVsOt46gsoJbYjMBQm5yKwHQqS6CaY91m8,3360
+phylogenie/treesimulator/events/mutations.py,sha256=fGxxeXhiFrMAaS-s_vd42aVBZ6XQiRjpW9uJUATNk1E,3752
 phylogenie/treesimulator/features.py,sha256=nNfv0kPudAZnrnufIVuFCR_VEJSP-6fZ2L-VG5_-Fqo,1366
 phylogenie/treesimulator/gillespie.py,sha256=_O1TqK-LOjH6k8ZSos-CbVQ7wJZPzf1Awv_1DXt0NYE,5730
 phylogenie/treesimulator/model.py,sha256=lhDwmBFQ8Qh8qVGZPgED0vehtPC3DE7_CgCV_8rPB-A,4641
 phylogenie/typeguards.py,sha256=JtqmbEWJZBRHbWgCvcl6nrWm3VcBfzRbklbTBYHItn0,1325
 phylogenie/typings.py,sha256=p694PBe_tk25A6N8vGGWxuqoDtt3nHFUsIYJrwR_76Y,494
 phylogenie/utils.py,sha256=ehVk_2kvjW8Q_EyM2kxBPHYiK-KlPmZQx7JeVN6Fh-E,5419
-phylogenie-2.1.26.dist-info/LICENSE.txt,sha256=NUrDqElK-eD3I0WqC004CJsy6cs0JgsAoebDv_42-pw,1071
-phylogenie-2.1.26.dist-info/METADATA,sha256=0-uiGyVuHCiWhvGtM93-Uq0Q9p5frdy6ppzgP5CDwzc,5477
-phylogenie-2.1.26.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-phylogenie-2.1.26.dist-info/entry_points.txt,sha256=Rt6_usN0FkBX1ZfiqCirjMN9FKOgFLG8rydcQ8kugeE,51
-phylogenie-2.1.26.dist-info/RECORD,,
+phylogenie-2.1.27.dist-info/LICENSE.txt,sha256=NUrDqElK-eD3I0WqC004CJsy6cs0JgsAoebDv_42-pw,1071
+phylogenie-2.1.27.dist-info/METADATA,sha256=UUZTU3YJ5yrAzQZ8EiiVsTmhuZzKNiGU_roQazkf3Mo,5477
+phylogenie-2.1.27.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+phylogenie-2.1.27.dist-info/entry_points.txt,sha256=Rt6_usN0FkBX1ZfiqCirjMN9FKOgFLG8rydcQ8kugeE,51
+phylogenie-2.1.27.dist-info/RECORD,,