forked from keon/algorithms
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathmarkov_chain.py
More file actions
73 lines (55 loc) · 1.81 KB
/
markov_chain.py
File metadata and controls
73 lines (55 loc) · 1.81 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
"""
Markov Chain
Provides utilities for stepping through and iterating a discrete Markov chain
described as a dictionary of transition probabilities.
Reference: https://en.wikipedia.org/wiki/Markov_chain
Complexity:
Time: O(S) per step, where S is the number of states
Space: O(S)
"""
from __future__ import annotations
import random
from collections.abc import Iterator
from typing import Any
def _choose_state(state_map: dict[Any, float]) -> Any | None:
"""Choose the next state randomly according to *state_map*.
Args:
state_map: Mapping of state to its transition probability.
Returns:
The selected state, or None if probabilities don't sum to 1.
"""
choice = random.random()
probability_reached = 0.0
for state, probability in state_map.items():
probability_reached += probability
if probability_reached > choice:
return state
return None
def next_state(chain: dict[Any, dict[Any, float]], current_state: Any) -> Any:
"""Return the next state given a Markov chain and current state.
Args:
chain: Markov chain as ``{state: {next_state: probability}}``.
current_state: The current state.
Returns:
The randomly chosen next state.
Examples:
>>> c = {'A': {'A': 1.0}}
>>> next_state(c, 'A')
'A'
"""
next_state_map = chain.get(current_state)
return _choose_state(next_state_map)
def iterating_markov_chain(
chain: dict[Any, dict[Any, float]],
state: Any,
) -> Iterator[Any]:
"""Yield an infinite sequence of states from a Markov chain.
Args:
chain: Markov chain transition dictionary.
state: Initial state.
Yields:
Successive states of the chain.
"""
while True:
state = next_state(chain, state)
yield state