Dynamical System Multilevel Model Notebook

.github/workflows/lint.yml

@@ -2,9 +2,9 @@ name: Lint
 
 on:
   push:
-    branches: [ master ]
+    branches: [ master, staging-dynamic ]
   pull_request:
-    branches: [ master ]
+    branches: [ master, staging-dynamic ]
 
 jobs:
   build:
@@ -32,7 +32,7 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install .[test]
+          pip install .[test,dynamical]
 
       - name: Lint
         run: ./scripts/lint.sh

.github/workflows/test.yml

@@ -2,9 +2,9 @@ name: Test
 
 on:
   push:
-    branches: [ master ]
+    branches: [ master, staging-dynamic ]
   pull_request:
-    branches: [ master ]
+    branches: [ master, staging-dynamic ]
 
 jobs:
   build:
@@ -46,7 +46,7 @@ jobs:
         run: |
           sudo apt install -y pandoc
           python -m pip install --upgrade pip
-          pip install .[test]
+          pip install .[test,dynamical]
 
       - name: Test
         shell: bash

chirho/dynamical/__init__.py

@@ -0,0 +1 @@
+from . import internals  # noqa: F401

chirho/dynamical/handlers/__init__.py

@@ -0,0 +1,12 @@
+from ..internals.solver import Solver  # noqa: F401
+from .event_loop import InterruptionEventLoop  # noqa: F401
+from .interruption import (  # noqa: F401
+    DynamicInterruption,
+    DynamicIntervention,
+    Interruption,
+    StaticBatchObservation,
+    StaticInterruption,
+    StaticIntervention,
+    StaticObservation,
+)
+from .trajectory import LogTrajectory  # noqa: F401

chirho/dynamical/handlers/event_loop.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+from typing import Generic, TypeVar
+
+import pyro
+
+from chirho.dynamical.handlers.interruption import Interruption
+from chirho.dynamical.internals.solver import (
+    apply_interruptions,
+    get_solver,
+    simulate_to_interruption,
+)
+
+S = TypeVar("S")
+T = TypeVar("T")
+
+
+class InterruptionEventLoop(Generic[T], pyro.poutine.messenger.Messenger):
+    def _pyro_simulate(self, msg) -> None:
+        dynamics, state, start_time, end_time = msg["args"]
+        if msg["kwargs"].get("solver", None) is not None:
+            solver = msg["kwargs"]["solver"]
+        else:
+            solver = get_solver()
+
+        # Simulate through the timespan, stopping at each interruption. This gives e.g. intervention handlers
+        #  a chance to modify the state and/or dynamics before the next span is simulated.
+        while start_time < end_time:
+            with pyro.poutine.messenger.block_messengers(
+                lambda m: m is self or (isinstance(m, Interruption) and m.used)
+            ):
+                state, terminal_interruptions, start_time = simulate_to_interruption(
+                    solver,
+                    dynamics,
+                    state,
+                    start_time,
+                    end_time,
+                )
+                for h in terminal_interruptions:
+                    h.used = True
+
+            with pyro.poutine.messenger.block_messengers(
+                lambda m: isinstance(m, Interruption)
+                and m not in terminal_interruptions
+            ):
+                dynamics, state = apply_interruptions(dynamics, state)
+
+        msg["value"] = state
+        msg["stop"] = True
+        msg["done"] = True
+        msg["in_SEL"] = True

chirho/dynamical/handlers/interruption.py

@@ -0,0 +1,177 @@
+import numbers
+import warnings
+from typing import Callable, Generic, Optional, TypeVar, Union
+
+import pyro
+import torch
+
+from chirho.dynamical.handlers.trajectory import LogTrajectory
+from chirho.dynamical.ops import State, get_keys
+from chirho.indexed.ops import get_index_plates, indices_of
+from chirho.interventional.ops import Intervention, intervene
+from chirho.observational.ops import Observation, observe
+
+R = Union[numbers.Real, torch.Tensor]
+S = TypeVar("S")
+T = TypeVar("T")
+
+
+class Interruption(pyro.poutine.messenger.Messenger):
+    used: bool
+
+    def __enter__(self):
+        self.used = False
+        return super().__enter__()
+
+    def _pyro_simulate_to_interruption(self, msg) -> None:
+        raise NotImplementedError("shouldn't be here!")
+
+
+class StaticInterruption(Interruption):
+    time: R
+
+    def __init__(self, time: R):
+        self.time = torch.as_tensor(time)  # TODO enforce this where it is needed
+        super().__init__()
+
+    def _pyro_simulate_to_interruption(self, msg) -> None:
+        _, _, _, start_time, end_time = msg["args"]
+
+        if start_time < self.time < end_time:
+            next_static_interruption: Optional[StaticInterruption] = msg["kwargs"].get(
+                "next_static_interruption", None
+            )
+
+            # Usurp the next static interruption if this one occurs earlier.
+            if (
+                next_static_interruption is None
+                or self.time < next_static_interruption.time
+            ):
+                msg["kwargs"]["next_static_interruption"] = self
+        elif self.time >= end_time:
+            warnings.warn(
+                f"{StaticInterruption.__name__} time {self.time} occurred after the end of the timespan "
+                f"{end_time}. This interruption will have no effect.",
+                UserWarning,
+            )
+
+
+class DynamicInterruption(Generic[T], Interruption):
+    """
+    :param event_f: An event trigger function that approaches and returns 0.0 when the event should be triggered.
+        This can be designed to trigger when the current state is "close enough" to some trigger state, or when an
+        element of the state exceeds some threshold, etc. It takes both the current time and current state.
+    """
+
+    def __init__(self, event_f: Callable[[R, State[T]], R]):
+        self.event_f = event_f
+        super().__init__()
+
+    def _pyro_simulate_to_interruption(self, msg) -> None:
+        msg["kwargs"].setdefault("dynamic_interruptions", []).append(self)
+
+
+class _InterventionMixin(Generic[T]):
+    """
+    We use this to provide the same functionality to both StaticIntervention and the DynamicIntervention,
+     while allowing DynamicIntervention to not inherit StaticInterruption functionality.
+    """
+
+    intervention: Intervention[State[T]]
+
+    def _pyro_apply_interruptions(self, msg) -> None:
+        dynamics, initial_state = msg["args"]
+        msg["args"] = (dynamics, intervene(initial_state, self.intervention))
+
+
+class _PointObservationMixin(Generic[T]):
+    observation: Observation[State[T]]
+    time: R
+
+    def _pyro_apply_interruptions(self, msg) -> None:
+        dynamics = msg["args"][0]
+        state: State[T] = msg["args"][1]
+        msg["value"] = (dynamics, observe(state, self.observation))
+
+    def _pyro_sample(self, msg):
+        # modify observed site names to handle multiple time points
+        msg["name"] = msg["name"] + "_" + str(torch.as_tensor(self.time).item())
+
+
+class StaticObservation(Generic[T], StaticInterruption, _PointObservationMixin[T]):
+    def __init__(
+        self,
+        time: R,
+        observation: Observation[State[T]],
+        *,
+        eps: float = 1e-6,
+    ):
+        self.observation = observation
+        # Add a small amount of time to the observation time to ensure that
+        # the observation occurs after the logging period.
+        super().__init__(time + eps)
+
+
+class StaticIntervention(Generic[T], StaticInterruption, _InterventionMixin[T]):
+    """
+    This effect handler interrupts a simulation at a given time, and
+    applies an intervention to the state at that time.
+
+    :param time: The time at which the intervention is applied.
+    :param intervention: The instantaneous intervention applied to the state when the event is triggered.
+    """
+
+    def __init__(self, time: R, intervention: Intervention[State[T]]):
+        self.intervention = intervention
+        super().__init__(time)
+
+
+class DynamicIntervention(Generic[T], DynamicInterruption, _InterventionMixin[T]):
+    """
+    This effect handler interrupts a simulation when the given dynamic event function returns 0.0, and
+    applies an intervention to the state at that time.
+
+    :param intervention: The instantaneous intervention applied to the state when the event is triggered.
+    """
+
+    def __init__(
+        self,
+        event_f: Callable[[R, State[T]], R],
+        intervention: Intervention[State[T]],
+    ):
+        self.intervention = intervention
+        super().__init__(event_f)
+
+
+class StaticBatchObservation(Generic[T], LogTrajectory[T]):
+    observation: Observation[State[T]]
+
+    def __init__(
+        self,
+        times: torch.Tensor,
+        observation: Observation[State[T]],
+        *,
+        eps: float = 1e-6,
+    ):
+        self.observation = observation
+        super().__init__(times, eps=eps)
+
+    def _pyro_post_simulate(self, msg) -> None:
+        super()._pyro_post_simulate(msg)
+
+        # This checks whether the simulate has already redirected in a InterruptionEventLoop.
+        # If so, we don't want to run the observation again.
+        if msg.setdefault("in_SEL", False):
+            return
+
+        # TODO remove this redundant check by fixing semantics of LogTrajectory and simulate
+        name_to_dim = {k: f.dim - 1 for k, f in get_index_plates().items()}
+        name_to_dim["__time"] = -1
+        len_traj = (
+            0
+            if not get_keys(self.trajectory)
+            else 1 + max(indices_of(self.trajectory, name_to_dim=name_to_dim)["__time"])
+        )
+
+        if len_traj == len(self.times):
+            msg["value"] = observe(self.trajectory, self.observation)

chirho/dynamical/handlers/solver.py

@@ -0,0 +1,16 @@
+from chirho.dynamical.internals.solver import Solver
+
+
+class TorchDiffEq(Solver):
+    def __init__(self, rtol=1e-7, atol=1e-9, method=None, options=None):
+        self.rtol = rtol
+        self.atol = atol
+        self.method = method
+        self.options = options
+        self.odeint_kwargs = {
+            "rtol": rtol,
+            "atol": atol,
+            "method": method,
+            "options": options,
+        }
+        super().__init__()

chirho/dynamical/handlers/trajectory.py

@@ -0,0 +1,70 @@
+import typing
+from typing import Generic, TypeVar
+
+import pyro
+import torch
+
+from chirho.dynamical.internals._utils import _squeeze_time_dim, append
+from chirho.dynamical.internals.solver import Solver, get_solver, simulate_trajectory
+from chirho.dynamical.ops import State
+from chirho.indexed.ops import IndexSet, gather, get_index_plates
+
+T = TypeVar("T")
+
+
+class LogTrajectory(Generic[T], pyro.poutine.messenger.Messenger):
+    trajectory: State[T]
+
+    def __init__(self, times: torch.Tensor, *, eps: float = 1e-6):
+        # Adding epsilon to the logging times to avoid collision issues with the logging times being exactly on the
+        #  boundaries of the simulation times. This is a hack, but it's a hack that should work for now.
+        self.times = times + eps
+
+        # Require that the times are sorted. This is required by the index masking we do below.
+        if not torch.all(self.times[1:] > self.times[:-1]):
+            raise ValueError("The passed times must be sorted.")
+
+        super().__init__()
+
+    def __enter__(self) -> "LogTrajectory[T]":
+        self.trajectory: State[T] = State()
+        return super().__enter__()
+
+    def _pyro_simulate(self, msg) -> None:
+        msg["done"] = True
+
+    def _pyro_post_simulate(self, msg) -> None:
+        # Turn a simulate that returns a state into a simulate that returns a trajectory at each of the logging_times
+        dynamics, initial_state, start_time, end_time = msg["args"]
+        if msg["kwargs"].get("solver", None) is not None:
+            solver = typing.cast(Solver, msg["kwargs"]["solver"])
+        else:
+            solver = get_solver()
+
+        filtered_timespan = self.times[
+            (self.times >= start_time) & (self.times <= end_time)
+        ]
+        timespan = torch.concat(
+            (start_time.unsqueeze(-1), filtered_timespan, end_time.unsqueeze(-1))
+        )
+
+        trajectory = simulate_trajectory(
+            solver,
+            dynamics,
+            initial_state,
+            timespan,
+        )
+
+        # TODO support dim != -1
+        idx_name = "__time"
+        name_to_dim = {k: f.dim - 1 for k, f in get_index_plates().items()}
+        name_to_dim[idx_name] = -1
+
+        if len(timespan) > 2:
+            part_idx = IndexSet(**{idx_name: set(range(1, len(timespan) - 1))})
+            new_part = gather(trajectory, part_idx, name_to_dim=name_to_dim)
+            self.trajectory: State[T] = append(self.trajectory, new_part)
+
+        final_idx = IndexSet(**{idx_name: {len(timespan) - 1}})
+        final_state = gather(trajectory, final_idx, name_to_dim=name_to_dim)
+        msg["value"] = _squeeze_time_dim(final_state)

chirho/dynamical/internals/__init__.py

@@ -0,0 +1,4 @@
+# Include only imports that are needed for registering dispatches.
+
+from . import _utils  # noqa: F401
+from . import backends  # noqa: F401