Fix time collision bug in LogTrajectory

[SamWitty]

This PR addresses #396 by adding the initial state to the trajectory if the start_time of simulate is equal to the first element of the (sorted) logging_times argument of the LogTrajectory handler. This works because the LogTrajectory handler excludes time collisions on the start_time and includes time collisions on the end_time of each simulate_point call, which was necessary to not "double include" interruption times that collided with an element in the logging_times and thus show up as the start_time and end_time of two simulate_point calls.

In addition, this PR adds a slight modification of the test described in #396 to the test suite.

[SamWitty]

@eb8680 , do you think it makes sense to move this line into _pyro_simulate instead of __enter__. This would change the behavior from "concatenate multiple simulate calls trajectories together" to "store only the trajectory of the final simulate call". Or alternatively, we could give simulate an optional name argument if you want to store multiple trajectories.

[eb8680]

Yes, re-initializing self.trajectory in _pyro_simulate seems more correct than the current behavior.

[eb8680]

Can you add a test that checks the case where a StaticInterruption's time coincides with a logging time, and also rerun the dynamical_intro notebook from scratch to verify that it no longer fails because of this bug?

[SamWitty]

Can you add a test that checks the case where a StaticInterruption's time coincides with a logging time, and also rerun the dynamical_intro notebook from scratch to verify that it no longer fails because of this bug?

I believe this existing test covers the test case? https://github.com/BasisResearch/chirho/blob/sw-time-collision/tests/dynamical/test_log_trajectory.py#L42 Happy to add more if that doesn't cover what you're looking for.

I'll rerun the notebook in this PR.

Would you like me to make the change we discussed in this comment in this PR or separately?

@eb8680 , do you think it makes sense to move this line into _pyro_simulate instead of __enter__. This would change the behavior from "concatenate multiple simulate calls trajectories together" to "store only the trajectory of the final simulate call". Or alternatively, we could give simulate an optional name argument if you want to store multiple trajectories.

[SamWitty]

@eb8680 , I edited this PR to address my suggestion here: #397 (comment). It is now ready for review again. I have also strengthened the tests a bit.

To ensure that handlers commute after this PR we now have the following strict order enforced regardless of the order handlers are applied:
LogTrajectory - Handles simulate during preprocessing to address potential time collision.
Solver - Handles simulate during postprocessing to do basically everything.
BatchObservation - Handles simulate to condition on data using a continuation, which occurs after all postprocessing.

I reran the dynamical system notebook, confirming that it works with these changes (although, requiring a minor change to the name of a trace address in the final cell because of a previous PR). I'm excluding the revised run from this PR to avoid merge conflicts with #377, which I expect to be completed by EOD or early tomorrow. Now that #308 is merged, I'll add a follow up PR after #377 to add the dynamical systems notebook to CI.

[eb8680]

Making sure LogTrajectory, Solver and BatchObservation commute is the right idea, but it seems like there's a simpler way to achieve that:

[eb8680]

I don't think you should need to change StaticBatchObservation at all, other than adding a single super()._pyro_post_simulate call:

def _pyro_post_simulate(self, msg: dict) -> None:
    super()._pyro_post_simulate(msg)  # update self.trajectory
    self.trajectory = observe(self.trajectory, self.observation)

[eb8680]

I think this change is right, and is already enough to make LogTrajectory and Solver commutative. In particular, since append is associative and Solver doesn't change initial_state or start_time in its _pyro_simulate method, as long as no outside code needs to access self.trajectory before the simulate call finishes, it shouldn't matter if LogTrajectory's _pyro_simulate runs before or after Solver's.

[eb8680]

This move from _pyro_simulate to _pyro_post_simulate seems unnecessary, per my comment on LogTrajectory.

[SamWitty]

Making sure LogTrajectory, Solver and BatchObservation commute is the right idea, but it seems like there's a simpler way to achieve that:

I agree. The key challenge here is actually in the assignment of self.trajectory: State[T] = State() in _pyro_simulate of LogTrajectory. Even though the append is associative, reseting self.trajectory must happen before the Solver's _pyro_simulate is called, or else it'll be entirely cleared in that line.

I'll try and find a workaround to this problem that's a bit less verbose and touches fewer pieces.

[eb8680]

The key challenge here is actually in the assignment of self.trajectory: State[T] = State() in _pyro_simulate of LogTrajectory. Even though the append is associative, reseting self.trajectory must happen before the Solver's _pyro_simulate is called, or else it'll be entirely cleared in that line.

Hmm, what about adding an extra boolean state variable to LogTrajectory that gets set in _pyro_post_simulate to track whether simulate has run? It's kind of clunky but at least it's isolated to LogTrajectory.

def __init__(self, ...):
    self._needs_reset: bool = False
    self.trajectory = State()
    ...

def __enter__(self):
    self._needs_reset = False
    self.trajectory = State()
    ...

def _pyro_simulate(self, msg: dict) -> None:
    if self._needs_reset:
        self.trajectory = State()
        self._needs_reset = False
    ...

def _pyro_post_simulate(self, msg: dict) -> None:
    if not self._needs_reset:
        self._needs_reset = True
    ...

[SamWitty]

@eb8680 , I've made this PR a bit simpler, but in the process it exposed an error that was occurring with a test using MultiWorldCounterfactual. I've decided to skip that test here rather than address it in this PR, as I believe doing so will require a solution to #379 . That's honestly a bit speculative though.

The (now skipped) test fails because dt.trajectory is never instantiated, whereas it was previously instantiated on __enter__ but never modified. As the test enumerated through the empty set of keys, it previously passed even though it shouldn't have.

[SamWitty]

The key challenge here is actually in the assignment of self.trajectory: State[T] = State() in _pyro_simulate of LogTrajectory. Even though the append is associative, reseting self.trajectory must happen before the Solver's _pyro_simulate is called, or else it'll be entirely cleared in that line.

Hmm, what about adding an extra boolean state variable to LogTrajectory that gets set in _pyro_post_simulate to track whether simulate has run? It's kind of clunky but at least it's isolated to LogTrajectory.

def __init__(self, ...):
    self._needs_reset: bool = False
    self.trajectory = State()
    ...

def __enter__(self):
    self._needs_reset = False
    self.trajectory = State()
    ...

def _pyro_simulate(self, msg: dict) -> None:
    if self._needs_reset:
        self.trajectory = State()
        self._needs_reset = False
    ...

def _pyro_post_simulate(self, msg: dict) -> None:
    if not self._needs_reset:
        self._needs_reset = True
    ...

My solution (that I finished before seeing this) is in the spirit of this suggestion, but a bit less explicit and a bit more concise. I'm really happy either way.

[eb8680]

LGTM, thanks for fixing this!