aboutsummaryrefslogtreecommitdiff
path: root/kennel/engine/components/markov_transition_state.py
blob: cec8223aa17dd7d5a6d611981ad5d10f8ab594c4 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
from kennel.engine.components.component import Component, ComponentType


class MarkovTransitionState(Component):
    def __init__(
        self,
        state_names: dict[int, str],
        initial_state_vector: list[float],
        transition_matrix: list[list[float]],
    ):
        # TODO: Poll rate per state?
        # TODO: State being an enum instead of a vector, just choose max and map
        self.state_names = state_names
        self.state = initial_state_vector
        self.transition_matrix = transition_matrix

        super().__init__(ComponentType.MARKOV)

    def get_max_state_name(self, state_vector: list[float]):
        max_val = max(state_vector)
        return self.state_names[state_vector.index(max_val)]

    def to_dict(self):
        return {"state": self.get_max_state_name(self.state)}