Skip to content

composite_expr

MarkovIVP

Bases: CompositeExpr

Uses scipy.integrate.solve_ivp to numerically find time evolution of a markov process given a transition rate matrix.

Returned shape is ,

Source code in slimfit/v2/composite_expr.py
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
class MarkovIVP(CompositeExpr):
    """Uses scipy.integrate.solve_ivp to numerically find time evolution of a markov process
        given a transition rate matrix.

    Returned shape is <states>, <datapoints>

    """

    def __init__(
        self,
        t: sp.Symbol | sp.Expr | Expr,
        trs_matrix: sp.Matrix | Expr,
        y0: sp.Matrix | Expr,
        domain: Optional[tuple[float, float]] = None,
        **ivp_kwargs,
    ):
        expr = as_expr({"t": t, "trs_matrix": trs_matrix, "y0": y0})
        super().__init__(expr)

        ivp_defaults = {"method": "Radau"}
        self.ivp_defaults = ivp_defaults | ivp_kwargs
        self.domain = domain

    def __call__(self, **kwargs):
        components = super().__call__(**kwargs)

        # if `self['t']` does not depend on any parameters; domain can be precomputed and
        # does not have to be determined for every call
        # although its every fast to do so
        domain = self.domain or self.get_domain(components["t"])
        sol = solve_ivp(
            self.grad_func,
            domain,
            y0=components["y0"].squeeze(),
            t_eval=components["t"],
            args=(components["trs_matrix"],),
            **self.ivp_defaults,
        )

        return sol.y

    def get_domain(self, arr: np.ndarray) -> tuple[float, float]:
        # padding?
        return arr[0], arr[-1]

    @staticmethod
    def grad_func(t, y, trs_matrix):
        return trs_matrix @ y