phylogenie 2.0.4__py3-none-any.whl → 2.0.6__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- phylogenie/generators/alisim.py +4 -3
- phylogenie/generators/configs.py +15 -19
- phylogenie/generators/dataset.py +3 -4
- phylogenie/generators/factories.py +9 -9
- phylogenie/generators/trees.py +62 -48
- phylogenie/generators/typeguards.py +9 -13
- phylogenie/main.py +4 -0
- phylogenie/skyline/matrix.py +3 -3
- phylogenie/skyline/parameter.py +4 -4
- phylogenie/skyline/vector.py +7 -9
- phylogenie/tree.py +19 -10
- phylogenie/treesimulator/__init__.py +2 -0
- phylogenie/treesimulator/events.py +89 -73
- phylogenie/treesimulator/gillespie.py +29 -12
- phylogenie/treesimulator/model.py +113 -59
- phylogenie/typeguards.py +10 -10
- {phylogenie-2.0.4.dist-info → phylogenie-2.0.6.dist-info}/METADATA +2 -2
- phylogenie-2.0.6.dist-info/RECORD +28 -0
- phylogenie-2.0.4.dist-info/RECORD +0 -28
- {phylogenie-2.0.4.dist-info → phylogenie-2.0.6.dist-info}/LICENSE.txt +0 -0
- {phylogenie-2.0.4.dist-info → phylogenie-2.0.6.dist-info}/WHEEL +0 -0
- {phylogenie-2.0.4.dist-info → phylogenie-2.0.6.dist-info}/entry_points.txt +0 -0
|
@@ -1,95 +1,84 @@
|
|
|
1
1
|
from abc import ABC, abstractmethod
|
|
2
|
+
from collections.abc import Sequence
|
|
3
|
+
from dataclasses import dataclass
|
|
2
4
|
|
|
3
|
-
|
|
5
|
+
import numpy as np
|
|
4
6
|
|
|
5
7
|
from phylogenie.skyline import (
|
|
6
8
|
SkylineMatrixCoercible,
|
|
9
|
+
SkylineParameter,
|
|
7
10
|
SkylineParameterLike,
|
|
8
11
|
SkylineVectorCoercible,
|
|
9
12
|
skyline_matrix,
|
|
10
13
|
skyline_parameter,
|
|
11
14
|
skyline_vector,
|
|
12
15
|
)
|
|
13
|
-
from phylogenie.treesimulator.model import Model
|
|
16
|
+
from phylogenie.treesimulator.model import Model, get_CT_state
|
|
14
17
|
|
|
15
18
|
INFECTIOUS_STATE = "I"
|
|
16
19
|
EXPOSED_STATE = "E"
|
|
17
20
|
SUPERSPREADER_STATE = "S"
|
|
18
21
|
|
|
19
22
|
|
|
23
|
+
@dataclass
|
|
20
24
|
class Event(ABC):
|
|
21
|
-
|
|
22
|
-
|
|
23
|
-
rate: SkylineParameterLike,
|
|
24
|
-
state: str | None = None,
|
|
25
|
-
):
|
|
26
|
-
self.rate = skyline_parameter(rate)
|
|
27
|
-
self.state = state
|
|
25
|
+
rate: SkylineParameter
|
|
26
|
+
state: str
|
|
28
27
|
|
|
29
28
|
def get_propensity(self, model: Model, time: float) -> float:
|
|
30
|
-
|
|
29
|
+
n_individuals = model.count_individuals(self.state)
|
|
30
|
+
rate = self.rate.get_value_at_time(time)
|
|
31
|
+
if rate == np.inf and not n_individuals:
|
|
32
|
+
return 0
|
|
33
|
+
return rate * n_individuals
|
|
31
34
|
|
|
32
35
|
@abstractmethod
|
|
33
|
-
def apply(self,
|
|
36
|
+
def apply(self, model: Model, time: float) -> None: ...
|
|
34
37
|
|
|
35
38
|
|
|
39
|
+
@dataclass
|
|
36
40
|
class BirthEvent(Event):
|
|
37
|
-
|
|
38
|
-
self,
|
|
39
|
-
rate: SkylineParameterLike,
|
|
40
|
-
state: str | None = None,
|
|
41
|
-
child_state: str | None = None,
|
|
42
|
-
):
|
|
43
|
-
super().__init__(rate, state)
|
|
44
|
-
self.child_state = child_state
|
|
41
|
+
child_state: str
|
|
45
42
|
|
|
46
|
-
def apply(self,
|
|
47
|
-
|
|
48
|
-
model.
|
|
43
|
+
def apply(self, model: Model, time: float) -> None:
|
|
44
|
+
individual = model.get_random_individual(self.state)
|
|
45
|
+
model.birth_from(individual, self.child_state, time)
|
|
49
46
|
|
|
50
47
|
|
|
51
48
|
class DeathEvent(Event):
|
|
52
|
-
def apply(self,
|
|
53
|
-
|
|
54
|
-
model.remove(
|
|
49
|
+
def apply(self, model: Model, time: float) -> None:
|
|
50
|
+
individual = model.get_random_individual(self.state)
|
|
51
|
+
model.remove(individual, time)
|
|
55
52
|
|
|
56
53
|
|
|
54
|
+
@dataclass
|
|
57
55
|
class MigrationEvent(Event):
|
|
58
|
-
|
|
59
|
-
super().__init__(rate, state)
|
|
60
|
-
self.target_state = target_state
|
|
56
|
+
target_state: str
|
|
61
57
|
|
|
62
|
-
def apply(self,
|
|
63
|
-
|
|
64
|
-
model.
|
|
58
|
+
def apply(self, model: Model, time: float) -> None:
|
|
59
|
+
individual = model.get_random_individual(self.state)
|
|
60
|
+
model.migrate(individual, self.target_state, time)
|
|
65
61
|
|
|
66
62
|
|
|
63
|
+
@dataclass
|
|
67
64
|
class SamplingEvent(Event):
|
|
68
|
-
|
|
69
|
-
self,
|
|
70
|
-
rate: SkylineParameterLike,
|
|
71
|
-
state: str | None = None,
|
|
72
|
-
removal_probability: SkylineParameterLike = 0,
|
|
73
|
-
):
|
|
74
|
-
super().__init__(rate, state)
|
|
75
|
-
self.removal_probability = skyline_parameter(removal_probability)
|
|
65
|
+
removal_probability: SkylineParameter
|
|
76
66
|
|
|
77
|
-
def apply(self,
|
|
78
|
-
|
|
79
|
-
|
|
80
|
-
model.sample(node, time, remove)
|
|
67
|
+
def apply(self, model: Model, time: float) -> None:
|
|
68
|
+
individual = model.get_random_individual(self.state)
|
|
69
|
+
model.sample(individual, time, self.removal_probability.get_value_at_time(time))
|
|
81
70
|
|
|
82
71
|
|
|
83
72
|
def get_canonical_events(
|
|
73
|
+
states: Sequence[str],
|
|
84
74
|
sampling_rates: SkylineVectorCoercible,
|
|
85
75
|
birth_rates: SkylineVectorCoercible = 0,
|
|
86
76
|
death_rates: SkylineVectorCoercible = 0,
|
|
87
77
|
removal_probabilities: SkylineVectorCoercible = 0,
|
|
88
78
|
migration_rates: SkylineMatrixCoercible | None = None,
|
|
89
79
|
birth_rates_among_states: SkylineMatrixCoercible | None = None,
|
|
90
|
-
states: list[str] | None = None,
|
|
91
80
|
) -> list[Event]:
|
|
92
|
-
N =
|
|
81
|
+
N = len(states)
|
|
93
82
|
|
|
94
83
|
birth_rates = skyline_vector(birth_rates, N)
|
|
95
84
|
death_rates = skyline_vector(death_rates, N)
|
|
@@ -97,43 +86,38 @@ def get_canonical_events(
|
|
|
97
86
|
removal_probabilities = skyline_vector(removal_probabilities, N)
|
|
98
87
|
|
|
99
88
|
events: list[Event] = []
|
|
100
|
-
for i in
|
|
101
|
-
state = None if states is None else states[i]
|
|
89
|
+
for i, state in enumerate(states):
|
|
102
90
|
events.append(BirthEvent(birth_rates[i], state, state))
|
|
103
91
|
events.append(DeathEvent(death_rates[i], state))
|
|
104
92
|
events.append(SamplingEvent(sampling_rates[i], state, removal_probabilities[i]))
|
|
105
93
|
|
|
106
|
-
if
|
|
94
|
+
if migration_rates is not None:
|
|
107
95
|
migration_rates = skyline_matrix(migration_rates, N, N - 1)
|
|
108
96
|
for i, state in enumerate(states):
|
|
109
97
|
for j, other_state in enumerate([s for s in states if s != state]):
|
|
110
98
|
events.append(MigrationEvent(migration_rates[i, j], state, other_state))
|
|
111
|
-
elif migration_rates is not None:
|
|
112
|
-
raise ValueError(f"Migration rates require states to be provided.")
|
|
113
99
|
|
|
114
|
-
if
|
|
100
|
+
if birth_rates_among_states is not None:
|
|
115
101
|
birth_rates_among_states = skyline_matrix(birth_rates_among_states, N, N - 1)
|
|
116
102
|
for i, state in enumerate(states):
|
|
117
103
|
for j, other_state in enumerate([s for s in states if s != state]):
|
|
118
104
|
events.append(
|
|
119
105
|
BirthEvent(birth_rates_among_states[i, j], state, other_state)
|
|
120
106
|
)
|
|
121
|
-
elif birth_rates_among_states is not None:
|
|
122
|
-
raise ValueError(f"Birth rates among states require states to be provided.")
|
|
123
107
|
|
|
124
108
|
return [event for event in events if event.rate]
|
|
125
109
|
|
|
126
110
|
|
|
127
111
|
def get_epidemiological_events(
|
|
112
|
+
states: Sequence[str],
|
|
128
113
|
sampling_proportions: SkylineVectorCoercible = 1,
|
|
129
114
|
reproduction_numbers: SkylineVectorCoercible = 0,
|
|
130
115
|
become_uninfectious_rates: SkylineVectorCoercible = 0,
|
|
131
116
|
removal_probabilities: SkylineVectorCoercible = 1,
|
|
132
117
|
migration_rates: SkylineMatrixCoercible | None = None,
|
|
133
118
|
reproduction_numbers_among_states: SkylineMatrixCoercible | None = None,
|
|
134
|
-
states: list[str] | None = None,
|
|
135
119
|
) -> list[Event]:
|
|
136
|
-
N =
|
|
120
|
+
N = len(states)
|
|
137
121
|
|
|
138
122
|
reproduction_numbers = skyline_vector(reproduction_numbers, N)
|
|
139
123
|
become_uninfectious_rates = skyline_vector(become_uninfectious_rates, N)
|
|
@@ -143,16 +127,14 @@ def get_epidemiological_events(
|
|
|
143
127
|
birth_rates = reproduction_numbers * become_uninfectious_rates
|
|
144
128
|
sampling_rates = become_uninfectious_rates * sampling_proportions
|
|
145
129
|
death_rates = become_uninfectious_rates - removal_probabilities * sampling_rates
|
|
146
|
-
birth_rates_among_states =
|
|
147
|
-
|
|
148
|
-
raise ValueError(
|
|
149
|
-
f"Reproduction numbers among states require states to be provided."
|
|
150
|
-
)
|
|
151
|
-
elif reproduction_numbers_among_states is not None:
|
|
152
|
-
birth_rates_among_states = (
|
|
130
|
+
birth_rates_among_states = (
|
|
131
|
+
(
|
|
153
132
|
skyline_matrix(reproduction_numbers_among_states, N, N - 1)
|
|
154
133
|
* become_uninfectious_rates
|
|
155
134
|
)
|
|
135
|
+
if reproduction_numbers_among_states is not None
|
|
136
|
+
else None
|
|
137
|
+
)
|
|
156
138
|
|
|
157
139
|
return get_canonical_events(
|
|
158
140
|
states=states,
|
|
@@ -166,15 +148,15 @@ def get_epidemiological_events(
|
|
|
166
148
|
|
|
167
149
|
|
|
168
150
|
def get_FBD_events(
|
|
151
|
+
states: Sequence[str],
|
|
169
152
|
diversification: SkylineVectorCoercible = 0,
|
|
170
153
|
turnover: SkylineVectorCoercible = 0,
|
|
171
154
|
sampling_proportions: SkylineVectorCoercible = 1,
|
|
172
155
|
removal_probabilities: SkylineVectorCoercible = 0,
|
|
173
156
|
migration_rates: SkylineMatrixCoercible | None = None,
|
|
174
157
|
diversification_between_types: SkylineMatrixCoercible | None = None,
|
|
175
|
-
states: list[str] | None = None,
|
|
176
158
|
) -> list[Event]:
|
|
177
|
-
N =
|
|
159
|
+
N = len(states)
|
|
178
160
|
|
|
179
161
|
diversification = skyline_vector(diversification, N)
|
|
180
162
|
turnover = skyline_vector(turnover, N)
|
|
@@ -185,15 +167,11 @@ def get_FBD_events(
|
|
|
185
167
|
death_rates = turnover * birth_rates
|
|
186
168
|
sampling_rates_dividend = 1 - removal_probabilities * sampling_proportions
|
|
187
169
|
sampling_rates = sampling_proportions * death_rates / sampling_rates_dividend
|
|
188
|
-
birth_rates_among_states =
|
|
189
|
-
|
|
190
|
-
|
|
191
|
-
|
|
192
|
-
|
|
193
|
-
elif diversification_between_types is not None:
|
|
194
|
-
birth_rates_among_states = (
|
|
195
|
-
skyline_matrix(diversification_between_types, N, N - 1) + death_rates
|
|
196
|
-
)
|
|
170
|
+
birth_rates_among_states = (
|
|
171
|
+
(skyline_matrix(diversification_between_types, N, N - 1) + death_rates)
|
|
172
|
+
if diversification_between_types is not None
|
|
173
|
+
else None
|
|
174
|
+
)
|
|
197
175
|
|
|
198
176
|
return get_canonical_events(
|
|
199
177
|
states=states,
|
|
@@ -212,6 +190,7 @@ def get_BD_events(
|
|
|
212
190
|
sampling_proportion: SkylineParameterLike = 1,
|
|
213
191
|
) -> list[Event]:
|
|
214
192
|
return get_epidemiological_events(
|
|
193
|
+
states=[INFECTIOUS_STATE],
|
|
215
194
|
reproduction_numbers=reproduction_number,
|
|
216
195
|
become_uninfectious_rates=1 / infectious_period,
|
|
217
196
|
sampling_proportions=sampling_proportion,
|
|
@@ -253,3 +232,40 @@ def get_BDSS_events(
|
|
|
253
232
|
become_uninfectious_rates=1 / infectious_period,
|
|
254
233
|
sampling_proportions=sampling_proportion,
|
|
255
234
|
)
|
|
235
|
+
|
|
236
|
+
|
|
237
|
+
def get_contact_tracing_events(
|
|
238
|
+
events: Sequence[Event],
|
|
239
|
+
samplable_states_after_notification: Sequence[str] | None = None,
|
|
240
|
+
sampling_rate_after_notification: SkylineParameterLike = np.inf,
|
|
241
|
+
contacts_removal_probability: SkylineParameterLike = 1,
|
|
242
|
+
) -> list[Event]:
|
|
243
|
+
ct_events = list(events)
|
|
244
|
+
for event in events:
|
|
245
|
+
if isinstance(event, MigrationEvent):
|
|
246
|
+
ct_events.append(
|
|
247
|
+
MigrationEvent(
|
|
248
|
+
event.rate,
|
|
249
|
+
get_CT_state(event.state),
|
|
250
|
+
get_CT_state(event.target_state),
|
|
251
|
+
)
|
|
252
|
+
)
|
|
253
|
+
elif isinstance(event, BirthEvent):
|
|
254
|
+
ct_events.append(
|
|
255
|
+
BirthEvent(event.rate, get_CT_state(event.state), event.child_state)
|
|
256
|
+
)
|
|
257
|
+
|
|
258
|
+
for state in (
|
|
259
|
+
samplable_states_after_notification
|
|
260
|
+
if samplable_states_after_notification is not None
|
|
261
|
+
else [e.state for e in events]
|
|
262
|
+
):
|
|
263
|
+
ct_events.append(
|
|
264
|
+
SamplingEvent(
|
|
265
|
+
skyline_parameter(sampling_rate_after_notification),
|
|
266
|
+
get_CT_state(state),
|
|
267
|
+
skyline_parameter(contacts_removal_probability),
|
|
268
|
+
)
|
|
269
|
+
)
|
|
270
|
+
|
|
271
|
+
return ct_events
|
|
@@ -3,9 +3,10 @@ from collections.abc import Sequence
|
|
|
3
3
|
import numpy as np
|
|
4
4
|
from numpy.random import default_rng
|
|
5
5
|
|
|
6
|
+
from phylogenie.skyline import SkylineParameterLike
|
|
6
7
|
from phylogenie.tree import Tree
|
|
7
|
-
from phylogenie.treesimulator.events import Event
|
|
8
|
-
from phylogenie.treesimulator.model import Model
|
|
8
|
+
from phylogenie.treesimulator.events import Event, get_contact_tracing_events
|
|
9
|
+
from phylogenie.treesimulator.model import Model, is_CT_state
|
|
9
10
|
|
|
10
11
|
|
|
11
12
|
def simulate_tree(
|
|
@@ -15,6 +16,11 @@ def simulate_tree(
|
|
|
15
16
|
max_time: float = np.inf,
|
|
16
17
|
init_state: str | None = None,
|
|
17
18
|
sampling_probability_at_present: float = 0.0,
|
|
19
|
+
notification_probability: float = 0,
|
|
20
|
+
max_notified_contacts: int = 1,
|
|
21
|
+
samplable_states_after_notification: Sequence[str] | None = None,
|
|
22
|
+
sampling_rate_after_notification: SkylineParameterLike = np.inf,
|
|
23
|
+
contacts_removal_probability: SkylineParameterLike = 1,
|
|
18
24
|
max_tries: int | None = None,
|
|
19
25
|
seed: int | None = None,
|
|
20
26
|
) -> Tree | None:
|
|
@@ -23,15 +29,19 @@ def simulate_tree(
|
|
|
23
29
|
if max_tips is None and max_time == np.inf:
|
|
24
30
|
raise ValueError("Either max_tips or max_time must be specified.")
|
|
25
31
|
|
|
32
|
+
if notification_probability:
|
|
33
|
+
events = get_contact_tracing_events(
|
|
34
|
+
events,
|
|
35
|
+
samplable_states_after_notification,
|
|
36
|
+
sampling_rate_after_notification,
|
|
37
|
+
contacts_removal_probability,
|
|
38
|
+
)
|
|
39
|
+
|
|
26
40
|
n_tries = 0
|
|
27
|
-
|
|
28
|
-
init_state = (
|
|
29
|
-
init_state
|
|
30
|
-
if init_state is not None
|
|
31
|
-
else str(rng.choice(states)) if states else None
|
|
32
|
-
)
|
|
41
|
+
root_states = [e.state for e in events if not is_CT_state(e.state)]
|
|
33
42
|
while max_tries is None or n_tries < max_tries:
|
|
34
|
-
|
|
43
|
+
root_state = init_state if init_state is not None else rng.choice(root_states)
|
|
44
|
+
model = Model(root_state, max_notified_contacts, notification_probability, rng)
|
|
35
45
|
current_time = 0.0
|
|
36
46
|
change_times = sorted(set(t for e in events for t in e.rate.change_times))
|
|
37
47
|
next_change_time = change_times.pop(0) if change_times else np.inf
|
|
@@ -39,6 +49,13 @@ def simulate_tree(
|
|
|
39
49
|
|
|
40
50
|
while current_time < max_time and (n_tips is None or model.n_sampled < n_tips):
|
|
41
51
|
rates = [e.get_propensity(model, current_time) for e in events]
|
|
52
|
+
|
|
53
|
+
instantaneous_events = [e for e, r in zip(events, rates) if r == np.inf]
|
|
54
|
+
if instantaneous_events:
|
|
55
|
+
event = instantaneous_events[rng.integers(len(instantaneous_events))]
|
|
56
|
+
event.apply(model, current_time)
|
|
57
|
+
continue
|
|
58
|
+
|
|
42
59
|
if not any(rates):
|
|
43
60
|
break
|
|
44
61
|
|
|
@@ -52,11 +69,11 @@ def simulate_tree(
|
|
|
52
69
|
continue
|
|
53
70
|
|
|
54
71
|
event_idx = np.searchsorted(np.cumsum(rates) / sum(rates), rng.random())
|
|
55
|
-
events[int(event_idx)].apply(
|
|
72
|
+
events[int(event_idx)].apply(model, current_time)
|
|
56
73
|
|
|
57
|
-
for
|
|
74
|
+
for individual in model.get_population():
|
|
58
75
|
if rng.random() < sampling_probability_at_present:
|
|
59
|
-
model.sample(
|
|
76
|
+
model.sample(individual, current_time, 1)
|
|
60
77
|
|
|
61
78
|
if model.n_sampled >= min_tips and (
|
|
62
79
|
max_tips is None or model.n_sampled <= max_tips
|
|
@@ -1,71 +1,125 @@
|
|
|
1
1
|
from collections import defaultdict
|
|
2
|
+
from dataclasses import dataclass, field
|
|
3
|
+
from typing import ClassVar
|
|
2
4
|
|
|
3
5
|
from numpy.random import Generator, default_rng
|
|
4
6
|
|
|
5
7
|
from phylogenie.tree import Tree
|
|
6
8
|
|
|
9
|
+
CT_POSTFIX = "-CT"
|
|
7
10
|
|
|
8
|
-
class Model:
|
|
9
|
-
def __init__(self, init_state: str | None = None):
|
|
10
|
-
self._next_id = 0
|
|
11
|
-
self._n_sampled = 0
|
|
12
|
-
self._leaves: dict[str, Tree] = {}
|
|
13
|
-
self._leaf2state: dict[str, str | None] = {}
|
|
14
|
-
self._state2leaves: dict[str | None, set[str]] = defaultdict(set)
|
|
15
|
-
self._tree = self._get_new_node(init_state, None)
|
|
16
11
|
|
|
17
|
-
|
|
18
|
-
|
|
19
|
-
|
|
20
|
-
|
|
12
|
+
def get_CT_state(state: str) -> str:
|
|
13
|
+
return f"{state}{CT_POSTFIX}"
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
def is_CT_state(state: str) -> bool:
|
|
17
|
+
return state.endswith(CT_POSTFIX)
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
@dataclass
|
|
21
|
+
class Individual:
|
|
22
|
+
node: Tree
|
|
23
|
+
state: str
|
|
24
|
+
id: int = field(init=False)
|
|
25
|
+
_id_counter: ClassVar[int] = 0
|
|
26
|
+
|
|
27
|
+
def __post_init__(self):
|
|
28
|
+
Individual._id_counter += 1
|
|
29
|
+
self.id = Individual._id_counter
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
class Model:
|
|
33
|
+
def __init__(
|
|
34
|
+
self,
|
|
35
|
+
init_state: str,
|
|
36
|
+
max_notified_contacts: int = 1,
|
|
37
|
+
notification_probability: float = 0,
|
|
38
|
+
rng: int | Generator | None = None,
|
|
39
|
+
):
|
|
40
|
+
self._next_node_id = 0
|
|
41
|
+
self._population: dict[int, Individual] = {}
|
|
42
|
+
self._states: dict[str, set[int]] = defaultdict(set)
|
|
43
|
+
self._contacts: dict[int, list[Individual]] = defaultdict(list)
|
|
44
|
+
self._sampled: set[str] = set()
|
|
45
|
+
self._tree = self._get_new_individual(init_state).node
|
|
46
|
+
self._max_notified_contacts = max_notified_contacts
|
|
47
|
+
self._notification_probability = notification_probability
|
|
48
|
+
self._rng = rng if isinstance(rng, Generator) else default_rng(rng)
|
|
21
49
|
|
|
22
50
|
@property
|
|
23
51
|
def n_sampled(self) -> int:
|
|
24
|
-
return self.
|
|
25
|
-
|
|
26
|
-
def _get_new_node(self, state: str
|
|
27
|
-
|
|
28
|
-
|
|
29
|
-
|
|
30
|
-
|
|
31
|
-
|
|
32
|
-
|
|
33
|
-
|
|
34
|
-
|
|
35
|
-
|
|
36
|
-
|
|
37
|
-
self._leaf2state.pop(node_id, None)
|
|
38
|
-
self._leaves.pop(node_id)
|
|
39
|
-
|
|
40
|
-
def add_child(
|
|
41
|
-
self,
|
|
42
|
-
node_id: str,
|
|
43
|
-
time: float,
|
|
44
|
-
stem: bool,
|
|
45
|
-
state: str | None,
|
|
46
|
-
branch_length: float | None = None,
|
|
47
|
-
) -> None:
|
|
48
|
-
node = self._leaves[node_id]
|
|
52
|
+
return len(self._sampled)
|
|
53
|
+
|
|
54
|
+
def _get_new_node(self, state: str) -> Tree:
|
|
55
|
+
self._next_node_id += 1
|
|
56
|
+
return Tree(f"{self._next_node_id}|{state}")
|
|
57
|
+
|
|
58
|
+
def _get_new_individual(self, state: str) -> Individual:
|
|
59
|
+
individual = Individual(self._get_new_node(state), state)
|
|
60
|
+
self._population[individual.id] = individual
|
|
61
|
+
self._states[state].add(individual.id)
|
|
62
|
+
return individual
|
|
63
|
+
|
|
64
|
+
def _set_branch_length(self, node: Tree, time: float) -> None:
|
|
49
65
|
if node.branch_length is not None:
|
|
50
|
-
raise ValueError("
|
|
51
|
-
node.add_child(self._get_new_node(state, branch_length))
|
|
52
|
-
if stem:
|
|
53
|
-
node.add_child(self._get_new_node(self._leaf2state[node.id], None))
|
|
66
|
+
raise ValueError(f"Branch length of node {node.id} is already set.")
|
|
54
67
|
node.branch_length = (
|
|
55
68
|
time if node.parent is None else time - node.parent.get_time()
|
|
56
69
|
)
|
|
57
|
-
self.remove(node_id)
|
|
58
70
|
|
|
59
|
-
def
|
|
60
|
-
self.
|
|
61
|
-
self.
|
|
71
|
+
def _stem(self, individual: Individual, time: float) -> None:
|
|
72
|
+
self._set_branch_length(individual.node, time)
|
|
73
|
+
stem_node = self._get_new_node(individual.state)
|
|
74
|
+
individual.node.add_child(stem_node)
|
|
75
|
+
individual.node = stem_node
|
|
76
|
+
|
|
77
|
+
def remove(self, id: int, time: float) -> None:
|
|
78
|
+
individual = self._population[id]
|
|
79
|
+
self._set_branch_length(individual.node, time)
|
|
80
|
+
state = individual.state
|
|
81
|
+
self._population.pop(id)
|
|
82
|
+
self._states[state].remove(id)
|
|
83
|
+
|
|
84
|
+
def migrate(self, id: int, state: str, time: float) -> None:
|
|
85
|
+
individual = self._population[id]
|
|
86
|
+
self._states[individual.state].remove(id)
|
|
87
|
+
individual.state = state
|
|
88
|
+
self._states[state].add(id)
|
|
89
|
+
self._stem(individual, time)
|
|
90
|
+
|
|
91
|
+
def birth_from(self, id: int, state: str, time: float) -> None:
|
|
92
|
+
individual = self._population[id]
|
|
93
|
+
new_individual = self._get_new_individual(state)
|
|
94
|
+
individual.node.add_child(new_individual.node)
|
|
95
|
+
self._stem(individual, time)
|
|
96
|
+
self._contacts[id].append(new_individual)
|
|
97
|
+
self._contacts[new_individual.id].append(individual)
|
|
98
|
+
|
|
99
|
+
def sample(self, id: int, time: float, removal_probability: float) -> None:
|
|
100
|
+
individual = self._population[id]
|
|
101
|
+
if self._rng.random() < removal_probability:
|
|
102
|
+
self._sampled.add(individual.node.id)
|
|
103
|
+
self.remove(id, time)
|
|
104
|
+
else:
|
|
105
|
+
sample_node = self._get_new_node(individual.state)
|
|
106
|
+
sample_node.branch_length = 0.0
|
|
107
|
+
self._sampled.add(sample_node.id)
|
|
108
|
+
individual.node.add_child(sample_node)
|
|
109
|
+
self._stem(individual, time)
|
|
110
|
+
|
|
111
|
+
for contact in self._contacts[id][-self._max_notified_contacts :]:
|
|
112
|
+
if (
|
|
113
|
+
contact.id in self._population
|
|
114
|
+
and not is_CT_state(contact.state)
|
|
115
|
+
and self._rng.random() < self._notification_probability
|
|
116
|
+
):
|
|
117
|
+
self.migrate(contact.id, get_CT_state(contact.state), time)
|
|
62
118
|
|
|
63
119
|
def get_sampled_tree(self) -> Tree:
|
|
64
120
|
tree = self._tree.copy()
|
|
65
121
|
for node in list(tree.postorder_traversal()):
|
|
66
|
-
if node.
|
|
67
|
-
node.branch_length > 0 and not node.children
|
|
68
|
-
):
|
|
122
|
+
if node.id not in self._sampled and not node.children:
|
|
69
123
|
if node.parent is None:
|
|
70
124
|
raise ValueError("No samples in the tree.")
|
|
71
125
|
else:
|
|
@@ -83,18 +137,18 @@ class Model:
|
|
|
83
137
|
node.parent.children.remove(node)
|
|
84
138
|
return tree
|
|
85
139
|
|
|
86
|
-
def
|
|
87
|
-
self
|
|
88
|
-
|
|
89
|
-
|
|
140
|
+
def get_full_tree(self) -> Tree:
|
|
141
|
+
return self._tree.copy()
|
|
142
|
+
|
|
143
|
+
def get_random_individual(self, state: str | None = None) -> int:
|
|
90
144
|
if state is None:
|
|
91
|
-
return
|
|
92
|
-
return
|
|
145
|
+
return self._rng.choice(list(self._population))
|
|
146
|
+
return self._rng.choice(list(self._states[state]))
|
|
93
147
|
|
|
94
|
-
def
|
|
95
|
-
return list(self.
|
|
148
|
+
def get_population(self) -> list[int]:
|
|
149
|
+
return list(self._population)
|
|
96
150
|
|
|
97
|
-
def
|
|
151
|
+
def count_individuals(self, state: str | None = None) -> int:
|
|
98
152
|
if state is None:
|
|
99
|
-
return len(self.
|
|
100
|
-
return len(self.
|
|
153
|
+
return len(self._population)
|
|
154
|
+
return len(self._states[state])
|
phylogenie/typeguards.py
CHANGED
|
@@ -1,42 +1,42 @@
|
|
|
1
1
|
from collections.abc import Sequence
|
|
2
|
-
from typing import TypeGuard
|
|
2
|
+
from typing import Any, TypeGuard
|
|
3
3
|
|
|
4
4
|
import phylogenie.typings as pgt
|
|
5
5
|
|
|
6
6
|
|
|
7
|
-
def is_many(x:
|
|
7
|
+
def is_many(x: Any) -> TypeGuard[pgt.Many[Any]]:
|
|
8
8
|
return isinstance(x, Sequence) and not isinstance(x, str)
|
|
9
9
|
|
|
10
10
|
|
|
11
|
-
def is_many_scalars(x:
|
|
11
|
+
def is_many_scalars(x: Any) -> TypeGuard[pgt.Many[pgt.Scalar]]:
|
|
12
12
|
return is_many(x) and all(isinstance(i, pgt.Scalar) for i in x)
|
|
13
13
|
|
|
14
14
|
|
|
15
|
-
def is_many_ints(x:
|
|
15
|
+
def is_many_ints(x: Any) -> TypeGuard[pgt.Many[int]]:
|
|
16
16
|
return is_many(x) and all(isinstance(i, int) for i in x)
|
|
17
17
|
|
|
18
18
|
|
|
19
|
-
def is_one_or_many_scalars(x:
|
|
19
|
+
def is_one_or_many_scalars(x: Any) -> TypeGuard[pgt.OneOrManyScalars]:
|
|
20
20
|
return isinstance(x, pgt.Scalar) or is_many_scalars(x)
|
|
21
21
|
|
|
22
22
|
|
|
23
|
-
def is_many_one_or_many_scalars(x:
|
|
23
|
+
def is_many_one_or_many_scalars(x: Any) -> TypeGuard[pgt.Many[pgt.OneOrManyScalars]]:
|
|
24
24
|
return is_many(x) and all(is_one_or_many_scalars(i) for i in x)
|
|
25
25
|
|
|
26
26
|
|
|
27
|
-
def is_many_2D_scalars(x:
|
|
27
|
+
def is_many_2D_scalars(x: Any) -> TypeGuard[pgt.Many2DScalars]:
|
|
28
28
|
return is_many(x) and all(is_many_scalars(i) for i in x)
|
|
29
29
|
|
|
30
30
|
|
|
31
|
-
def is_one_or_many_2D_scalars(x:
|
|
31
|
+
def is_one_or_many_2D_scalars(x: Any) -> TypeGuard[pgt.OneOrMany2DScalars]:
|
|
32
32
|
return isinstance(x, pgt.Scalar) or is_many_2D_scalars(x)
|
|
33
33
|
|
|
34
34
|
|
|
35
35
|
def is_many_one_or_many_2D_scalars(
|
|
36
|
-
x:
|
|
36
|
+
x: Any,
|
|
37
37
|
) -> TypeGuard[pgt.Many[pgt.OneOrMany2DScalars]]:
|
|
38
38
|
return is_many(x) and all(is_one_or_many_2D_scalars(i) for i in x)
|
|
39
39
|
|
|
40
40
|
|
|
41
|
-
def is_many_3D_scalars(x:
|
|
41
|
+
def is_many_3D_scalars(x: Any) -> TypeGuard[pgt.Many3DScalars]:
|
|
42
42
|
return is_many(x) and all(is_many_2D_scalars(i) for i in x)
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.1
|
|
2
2
|
Name: phylogenie
|
|
3
|
-
Version: 2.0.
|
|
3
|
+
Version: 2.0.6
|
|
4
4
|
Summary: Generate phylogenetic datasets with minimal setup effort
|
|
5
5
|
Author: Gabriele Marino
|
|
6
6
|
Author-email: gabmarino.8601@gmail.com
|
|
@@ -45,7 +45,7 @@ Phylogenie comes packed with useful features, including:
|
|
|
45
45
|
Simply specify the number of cores to use, and Phylogenie handles multiprocessing automatically.
|
|
46
46
|
|
|
47
47
|
- **Pre-implemented parameterizations** 🎯
|
|
48
|
-
Include canonical, fossilized birth-death, epidemiological, birth-death with exposed-infectious (BDEI), birth-death with superspreading (BDSS), and
|
|
48
|
+
Include canonical, fossilized birth-death, epidemiological, birth-death with exposed-infectious (BDEI), birth-death with superspreading (BDSS), and contact tracing (CT).
|
|
49
49
|
|
|
50
50
|
- **Skyline parameter support** 🪜
|
|
51
51
|
Support for piece-wise constant parameters.
|
|
@@ -0,0 +1,28 @@
|
|
|
1
|
+
phylogenie/__init__.py,sha256=1w_0H9lg7hI3b-NLjKuzc34GbJJGyjLq9LrlogecTzI,1759
|
|
2
|
+
phylogenie/generators/__init__.py,sha256=zsOxy28-9j9alOQLIgrOAFfmM58NNHO_NEtW-KXQXAY,888
|
|
3
|
+
phylogenie/generators/alisim.py,sha256=dDqlSwLDbRE2u5SZlsq1mArobTBtuk0aeXY3m1N-bWA,2374
|
|
4
|
+
phylogenie/generators/configs.py,sha256=5ZWdKhRUjlNifw7QKXbooKV1fElqfCk_jBGxfcjh8do,969
|
|
5
|
+
phylogenie/generators/dataset.py,sha256=k6RYJpgxOL8a_yMq98WUF-dcJv8TwxaWnde0k13M4J0,2525
|
|
6
|
+
phylogenie/generators/factories.py,sha256=O8wqL-PvZps0Dq6mQa_PTi4vBvky5LkQIy1jjfOUm-4,6944
|
|
7
|
+
phylogenie/generators/trees.py,sha256=jukaVXGcPGzDBEYMGJ1MKqWt4XbAB5EEfuHXDpwKTqM,9173
|
|
8
|
+
phylogenie/generators/typeguards.py,sha256=Qph6ZnQ7wDMUNvB0VWQKlq42f8wkKOnM42cfMqhNov4,862
|
|
9
|
+
phylogenie/io.py,sha256=ZXlofnSh7FX5UJiP0svRHrTraMSNgKa1GiAv0bMz7jU,2854
|
|
10
|
+
phylogenie/main.py,sha256=4mhy3qmCPrGUm7QHAZjuLoR39eiFnhdw-C0mPh3lzxU,1264
|
|
11
|
+
phylogenie/msa.py,sha256=JDGyZUsAq6-m-SQjoCDjAkAZIxfgyl_PDIhdYn5HOow,2064
|
|
12
|
+
phylogenie/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
13
|
+
phylogenie/skyline/__init__.py,sha256=7pF4CUb4ZCLzNYJNhOjpuTOLTRhlK7L6ugfccNqjIGo,620
|
|
14
|
+
phylogenie/skyline/matrix.py,sha256=Gl8OgKjtieG0NwPYiPimKI36gefV8fm_OeorjdXxPTs,9146
|
|
15
|
+
phylogenie/skyline/parameter.py,sha256=EM9qlPt0JhMBy3TbztM0dj24BaGNEy8KWKdTObDKhbI,4644
|
|
16
|
+
phylogenie/skyline/vector.py,sha256=bJP7_FNX_Klt6wXqsyfj0KX3VNj6-dIhzCKSJuQcOV0,7115
|
|
17
|
+
phylogenie/tree.py,sha256=34gcxUoTGfj72EbIlpnrhWGnNFppUVjms3XEn1ZS3-g,1997
|
|
18
|
+
phylogenie/treesimulator/__init__.py,sha256=INPU9LrPdUmt3dYGzWDRoRKrPR9xENcHu44pJVUbyNA,525
|
|
19
|
+
phylogenie/treesimulator/events.py,sha256=X3_0U9qqMpYgh6-7TwQEnlUipANkHz6QTCXlm-qXFQk,9524
|
|
20
|
+
phylogenie/treesimulator/gillespie.py,sha256=Fn-PyVICx3pWtpHko7rf6omf_kqOkkpebSJy56oPKnQ,3216
|
|
21
|
+
phylogenie/treesimulator/model.py,sha256=XpzAicmg2O6K0Trk5YolH-B_HJZxoSauF2wZOMqp-Iw,5559
|
|
22
|
+
phylogenie/typeguards.py,sha256=JtqmbEWJZBRHbWgCvcl6nrWm3VcBfzRbklbTBYHItn0,1325
|
|
23
|
+
phylogenie/typings.py,sha256=O1X6lGKTjJ2YJz3ApQ-rYb_tEJNUIcHdUIeYlSM4s5o,500
|
|
24
|
+
phylogenie-2.0.6.dist-info/LICENSE.txt,sha256=NUrDqElK-eD3I0WqC004CJsy6cs0JgsAoebDv_42-pw,1071
|
|
25
|
+
phylogenie-2.0.6.dist-info/METADATA,sha256=AkoTSMmVcV2sFdsnGm-ldW_Oj25C83T5_1V2-nF8KUc,5472
|
|
26
|
+
phylogenie-2.0.6.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
|
|
27
|
+
phylogenie-2.0.6.dist-info/entry_points.txt,sha256=Rt6_usN0FkBX1ZfiqCirjMN9FKOgFLG8rydcQ8kugeE,51
|
|
28
|
+
phylogenie-2.0.6.dist-info/RECORD,,
|