Adding support for iterating through MDAnalysis trajectories

[TENeary]

Quick PR to add ability for MDAnalysisSimulation protocols to iterate through Universes.

[TENeary]

Adding a comment to ensure there's a note regarding the compatibility between the Omni runner and this implementation of the Simulation and to make sure this is the intended functionality.
Currently the play signal from the Omni runner uses the advance_by_seconds func from the simulation protocol.
The runner expects this to relate to the real world time elapsed between the previous sent frame but passing this as the change to simulation time doesn't make sense given simulation time/real world time.

Should Omni runner dt passed be altered or use a different method especially given all current protocols (this PR excluded) only advance to the next step regardless.

[Ragzouken]

I'm beginning to think this whole constructor doesn't make sense, and we should only have from_universe instead of proxying the Universe constructor.

I'd remove from_path and have the CLI just construct the universe itself before passing it in.

[Ragzouken]

Adding a comment to ensure there's a note regarding the compatibility between the Omni runner and this implementation of the Simulation and to make sure this is the intended functionality. Currently the play signal from the Omni runner uses the advance_by_seconds func from the simulation protocol. The runner expects this to relate to the real world time elapsed between the previous sent frame but passing this as the change to simulation time doesn't make sense given simulation time/real world time.

Should Omni runner dt passed be altered or use a different method especially given all current protocols (this PR excluded) only advance to the next step regardless.

As far as I remember, advance_by_seconds really only exists to make it possible to run the live recordings back in real time. Probably it would be better that each simulation just tracks that internally rather than expecting the server to tell it how time is passing--but exactly how/when frames are generated needs more thought it general.

For this PR it probably makes sense to just pick some sensible mapping factor of real world time to simulation time (possibly exposed as a property for convenience) and leave it at that. Recording/trajectory playback is underdesigned so I don't think there's much more to be done there in this PR.

[Ragzouken]

I updated the mdanalysis_trajectory tutorial notebook to use this method, but playback is either static or too slow to see anything--picking a default time factor or some other change that makes the playback in that notebook proceed as expected would be good to demonstrate playback is working (using the VR step controls manually works fine).

[TENeary]

I updated the mdanalysis_trajectory tutorial notebook to use this method, but playback is either static or too slow to see anything--picking a default time factor or some other change that makes the playback in that notebook proceed as expected would be good to demonstrate playback is working (using the VR step controls manually works fine).

If you're using the 3TIC (ose) simulation in the notebook for testing, the lack of an observed change is due to my previous comment regarding advancing by seconds in real world vs. simulation time. Assuming the runner interval is set at roughly 30 FPS the simulation will attempt to jump 33 frames - as the simulation only has 24, it just loops back to the beginning so you don't see any change.

[Ragzouken]

I've made an update based on trying to get the trajectory playback notebook to work reasonably: the simulation keeps track of simulation time elapsed during playback rather than only which frame is it currently on--this way even a small dt eventually adds up to enough time passing to advance between frames. This is similar to how the recording playback works. Advancing by a step snaps the simulation time to the first time where the next frame would be shown.

Do you see any issue with this? The notebook is the only in-repo example of using UniverseSimulation so it should at least match the use case there, but I don't know if there's another use case where this breaks things.

I put it on another branch so I don't overwrite anything: https://github.com/IRL2/nanover-server-py/compare/feature/mdanalysis_sim_variant

[TENeary]

The seek_to_time is a good idea for sure. A lot of the tracking is already done by MDA so I've simplified the implementation a little. It should work now also with variable speed (via playback_factor or time_to_step assuming you advance by report or step).
I've also update the CLI interface for mda to allow for passing topology and multiple trajectories. Currently it doesn't raise an exception if it can't create the Universe - I wasn't sure if we want to force an exist if any of the simulations fail to load or not.
The timestep management may be a little overengineered but let me what you think.

[Ragzouken]

With your changes it works as expected in the notebook, which is good, but he logic of how it works seems quite convoluted and difficult to follow.

The two ways simulations are advanced are advance_by_one_step and advance_by_seconds -- one for stepping and one for unpaused playback. Both of these set current_dt which takes priority over default_dt. So then what is the purpose of the time_to_step property?

StepManager has a next_ts method but it doesn't seem to be called by anything.

In if (next_ts := next(self._time_to_step)) is None: there is a lot going on, it would be better to break it up, but regardless is there ever a case where this is actually None anyway?

I didn't realise this before but iter(universe.trajectory) just returns universe.trajectory so _universe_iterator can be removed and replaced with universe.trajectory. Then you can see that the load method should just be empty, which makes sense because the Universe is already ready at the time of creation.

_next_frame has an optional reset argument but it's never passed by anything, and even if reset were False, the next invocation of the method is going to restart the iterator either way (next(universe.trajectory) after the StopIteration exception just restarts the iterator).

Because this is quite difficult to reason about and there seem to be a few unused or misleading parts, I'm going to switch us to the simplified version I made in the other branch. Let's use that as a base and we can discuss if it's its missing some important behavior from the previous version.

[TENeary]

With your changes it works as expected in the notebook, which is good, but he logic of how it works seems quite convoluted and difficult to follow.

Because this is quite difficult to reason about and there seem to be a few unused or misleading parts, I'm going to switch us to the simplified version I made in the other branch. Let's use that as a base and we can discuss if it's its missing some important behavior from the previous version.

I disagree, the new implementation makes a number of odd choices regarding naming that do not conform to the existing Protocol the class is derived from and a number of regressive edits handling edge cases that will need to be addressed again.
I conceed that StepTracker may be clunky but removing it entirely without recapturing the original aim of its implementation seems excessive to me.

The previous implementation aimed to enable 3 distinct functionality via the advance_by.. methods:

• advance_by_seconds maps real world time elapsing to sim time.
• advance_by_one_step step the simulation
• advance_to_next_report advance the simulation by user defined period of simulation time - controlled via StepManager (see notes below).

The new implenentation does not allow for this. I would argue playback mapping to real world seconds is a more niche usecase than simply advancing a fixed time step every frame. Further removing advance_to_next_report muddles the codebase in that every other Simulation type implements this method - rather its the only method they actually implement - so it not being consistent across all types is inconsistent.
I also don't understand the rationale for renaming the "advance" methods to "seek" in general. While I agree this is better reflects their actual function it goes against the current naming scheme from the existing protocols.
This also makes any sort of generic interface with using these classes significantly more difficult to implement.
If anything, it should be the implementation in this Simulation that sets the precedence for the first two methods as ASE and OpenMM are always ultimately calling the same advance_to_next_report.

The two ways simulations are advanced are advance_by_one_step and advance_by_seconds -- one for stepping and one for unpaused playback. Both of these set current_dt which takes priority over default_dt. So then what is the purpose of the time_to_step property?

The goal of the StepManager was to enable a method for user to define a default step size when using the advance_to_next_report - i.e. the default_dt. This parameter is accessed whenever current_dt is None.
This allows the time_to_step for a default step time and optionally arbitrary steps to be used simultaneously. While I am open to an alternative implementation, I do think there is use cases for this.
Refer to above for why I believe that renaming methods is a poor idea without it being systematic.

In if (next_ts := next(self._time_to_step)) is None: there is a lot going on, it would be better to break it up, but regardless is there ever a case where this is actually None anyway?

This can be None when the default time to step (default_dt) and current_dt are both None, giving a fallback of just advancing a single frame. The next_ts property is a means to grab the next intended time jump without exhausting the currently set current_dt which is converted to None after next() is called. While not called in the code it is useful for checking the intended time jump before it is done.

I didn't realise this before but iter(universe.trajectory) just returns universe.trajectory so _universe_iterator can be removed and replaced with universe.trajectory. Then you can see that the load method should just be empty, which makes sense because the Universe is already ready at the time of creation.

Happy to change iter(...) to just universe.trajectory. the argument against removing the iter statements is to prevent UB resulting frmo when a user still has a reference to the universe. They may move the trajectory frame and keeping helps to ensure the simulation is always set to the first timestep.
I don't see a reason to do a bunch of additional next_frame/frame_count calculations when all of this is already stored/managed by the simulation. To prevent UB wouldn't it make more sense to use MDA's internal data? Additionally, the only instance where it is actually necessary to track any time progressing is for the real world relating to fracitonal increases in simulation dt - i.e. essentailly advancing by a fraction of a frame.

_next_frame has an optional reset argument but it's never passed by anything, and even if reset were False, the next invocation of the method is going to restart the iterator either way (next(universe.trajectory) after the StopIteration exception just restarts the iterator).

The goal of reset in _next_frame is allow for cases where a user wants to halt a simulation playback at the end of its timecourse instead of running in perpetuity. Perhaps this functionality is better handled by a member variable so it can more easily set by the user given the current restrictions in the function signature of the Simulation Protocol. But I believe this functionality is important to retain.

Also, from what I understand of this current implementation, the seek_to_time also has (what I would consider) poor behaviour. If a user provides a time > simulation time the Simulation moves to some frame in the middle of the simulation. e.g. passing 130s into simulation with a total time of 100s would jump to the frame representing 30s.
With maybe the exception of mapping playback to seconds elapsed in the real world this doesn't really seem to be a sensible resolution. I would argue this should either return the last frame (or first, if looping) or raise some error.

[Ragzouken]

I addressed things individually, but the long and the short of it for me is still: the previous version was difficult to follow and the complexity obscured what seem to be a number of redundancies and bugs. The alternate implementation is two completely linear methods to implement the two required methods of the Simulation protocol. It drops the time_to_step property, but I don't see a use case for it that isn't covered by advance_by_seconds. Seeking out of bounds has perhaps odd behaviour, maybe it's worth accounting for that. I think the naming is fine in the context, but it's a minor change anyway.

I didn't realise this before but iter(universe.trajectory) just returns universe.trajectory so _universe_iterator can be removed and replaced with universe.trajectory. Then you can see that the load method should just be empty, which makes sense because the Universe is already ready at the time of creation.

Happy to change iter(...) to just universe.trajectory. the argument against removing the iter statements is to prevent UB resulting frmo when a user still has a reference to the universe. They may move the trajectory frame and keeping helps to ensure the simulation is always set to the first timestep. I don't see a reason to do a bunch of additional next_frame/frame_count calculations when all of this is already stored/managed by the simulation. To prevent UB wouldn't it make more sense to use MDA's internal data? Additionally, the only instance where it is actually necessary to track any time progressing is for the real world relating to fracitonal increases in simulation dt - i.e. essentailly advancing by a fraction of a frame.

My point here was that universe.trajectory.__iter__ returns itself, and storing it as a member gives the misleading impression that its an independent object with its own iteration state. This is not the case: iteration over MDAnalysis universes is singleton, and has side-effects on the universe/trajectory itself. I assume avoiding this is what you were refering to with UB (undefined behaviour?)?

_next_frame has an optional reset argument but it's never passed by anything, and even if reset were False, the next invocation of the method is going to restart the iterator either way (next(universe.trajectory) after the StopIteration exception just restarts the iterator).

The goal of reset in _next_frame is allow for cases where a user wants to halt a simulation playback at the end of its timecourse instead of running in perpetuity. Perhaps this functionality is better handled by a member variable so it can more easily set by the user given the current restrictions in the function signature of the Simulation Protocol. But I believe this functionality is important to retain.

My point here is that 1. _next_frame is an internal function and internally it is never called with reset=False so I don't see how it could ever be utilised, and 2. that it doesn't function as described (if I change the default to False it still loops) because MDAnalysis readers autorewind when they reach the end.

The existing recording playback PlaybackSimulation doesn't support non-looping playback and I'm not aware of it having been requested, so I feel its not an essential consideration.

With your changes it works as expected in the notebook, which is good, but he logic of how it works seems quite convoluted and difficult to follow.
Because this is quite difficult to reason about and there seem to be a few unused or misleading parts, I'm going to switch us to the simplified version I made in the other branch. Let's use that as a base and we can discuss if it's its missing some important behavior from the previous version.

I disagree, the new implementation makes a number of odd choices regarding naming that do not conform to the existing Protocol the class is derived from and a number of regressive edits handling edge cases that will need to be addressed again. I conceed that StepTracker may be clunky but removing it entirely without recapturing the original aim of its implementation seems excessive to me.

advance_to_next_report is not part of the Simulation interface, it's just a common element of OpenMMSimulation and ASESimulation, the two implementations of live playback. In those it's essentially a synonym of advance_by_one_step but with a different name to emphasize that the "step" of the Simulation protocol is not the same thing as a "step" in OpenMM or ASE--in this way advance_to_next_report is already a case of deviating from the existing names to provide more clarity.

Since the methods for jumping to particular timestamps/frames don't form part of the Simulation protocol (they wouldn't even really make sense for live dynamics), but are rather bespoke to this class, I don't see any reason for them to sacrifice clearer naming.

The previous implementation aimed to enable 3 distinct functionality via the advance_by.. methods:

• advance_by_seconds maps real world time elapsing to sim time.
• advance_by_one_step step the simulation
• advance_to_next_report advance the simulation by user defined period of simulation time - controlled via StepManager (see notes below).

The new implenentation does not allow for this. I would argue playback mapping to real world seconds is a more niche usecase than simply advancing a fixed time step every frame. Further removing advance_to_next_report muddles the codebase in that every other Simulation type implements this method - rather its the only method they actually implement - so it not being consistent across all types is inconsistent. I also don't understand the rationale for renaming the "advance" methods to "seek" in general. While I agree this is better reflects their actual function it goes against the current naming scheme from the existing protocols. This also makes any sort of generic interface with using these classes significantly more difficult to implement. If anything, it should be the implementation in this Simulation that sets the precedence for the first two methods as ASE and OpenMM are always ultimately calling the same advance_to_next_report.

The rationale for mapping time is just that since the trajectories can have different time resolutions and different frequency of frame reporting, it will lead to different rates of frame playback that can end up being way too slow or way to fast. The existing example for trajectory playback (https://github.com/IRL2/nanover-server-py/blob/main/tutorials/mdanalysis/mdanalysis_trajectory.ipynb) provided playback_fps for this purpose. If that's actually an extraneous consideration then I have no problem dropping it.

The two ways simulations are advanced are advance_by_one_step and advance_by_seconds -- one for stepping and one for unpaused playback. Both of these set current_dt which takes priority over default_dt. So then what is the purpose of the time_to_step property?

The goal of the StepManager was to enable a method for user to define a default step size when using the advance_to_next_report - i.e. the default_dt. This parameter is accessed whenever current_dt is None. This allows the time_to_step for a default step time and optionally arbitrary steps to be used simultaneously. While I am open to an alternative implementation, I do think there is use cases for this. Refer to above for why I believe that renaming methods is a poor idea without it being systematic.

In if (next_ts := next(self._time_to_step)) is None: there is a lot going on, it would be better to break it up, but regardless is there ever a case where this is actually None anyway?

This can be None when the default time to step (default_dt) and current_dt are both None, giving a fallback of just advancing a single frame. The next_ts property is a means to grab the next intended time jump without exhausting the currently set current_dt which is converted to None after next() is called. While not called in the code it is useful for checking the intended time jump before it is done.

I don't really understand the motivation for this functionality--who is it for? It only works when you call advance_to_next_report manually, so why not call advance_by_seconds with the desired alternate timestep?

In any case I think it obscures things 1. that it is a seperate class, and 2. that its presented as an iterator, when it encapsulates logic that is only used in exactly one place (_next_frame).

My intention with reworking it is just be explicit that there are two ways to manipulate this object--go to the next frame, or go to a point in the playback time--and then to make the implementations of those two operations completely clear about what they are doing: determine the current frame and move to the next one and update time, or update time and determine the resulting frame.

It's difficult to follow the original version where the internal state of playback is hidden inside StepManager and universe.trajectory, and you need to understand exactly how the MDAnalysis reader iterators work, trace the flow of step time to figure out which branch of _next_frame you're going to be in and why.

Similar, to me it's clearer to be explicit about what duration each loop is supposed to be by writing:

frame_count = len(self.universe.trajectory)
frame_length = self.universe.trajectory.dt
duration = frame_count * frame_length

Whereas, as far as I can tell, using universe.trajectory.totaltime has introduced a bug of always skipping the last frame in advance_by_seconds because totaltime is the interval between first and last frame, leaving no extra time afterwards before looping to actually display that frame.

Also, from what I understand of this current implementation, the seek_to_time also has (what I would consider) poor behaviour. If a user provides a time > simulation time the Simulation moves to some frame in the middle of the simulation. e.g. passing 130s into simulation with a total time of 100s would jump to the frame representing 30s. With maybe the exception of mapping playback to seconds elapsed in the real world this doesn't really seem to be a sensible resolution. I would argue this should either return the last frame (or first, if looping) or raise some error.

The intention here was just to avoid having to handle looping as a special case, but if the users find it confusing, it can be changed of course.

[TENeary]

For the minor stuff, e.g. frame jump computations - I don't have any super critical issues its mostly semantic issues or me attempting to future proof the code for future issues I see arising.
I expect that I didn't make my intentions particularly clear in the previous comment.
Regarding these small issues:

• I expect I wasn't clear enough - I have no issues dropping the universe_iterator member - especially as all functionality can be handled by manipulating the underlying .trajectory.
• seek_to_next_frame is overengineered, in my opinion. It can simply be resolved with next(universe.trajectory) + some minor work to handle reaching the end of the trajectory. There is no need to compute frame indexes etc.
• The iter point I referred to is for the load method specifically. I intended to mean that the simulation should always be reset on load - i.e. it sets the expectation that anytime a simulation is loaded it will always start from the beginning.
• ...totaltime bug is a minor off by 1 error resolved with ...totaltime + universe.trajectory.dt
• I wasn't aware of the reset bug in _next_frame but it can simply be resolved by jumping to the last timestep after a StopIteration (e.g. universe.trajectory[-1]).

That said, I still believe there are 2 key issues to be resolved:

Regarding the 3 seperate functionalities:
I strongly disagree that the advance by a set number of simulation time (advance_to_next_report) in my previous comment is unneeded. I would expect advance_by_seconds in a class handling MDA simulations to advance the simulation to the frame corresponding to the the time dt seconds ahead. I, and I expect for most other people, would see this as the intuitive option.
From the function name, you would expect the simulation to advance by some number of (nano)seconds - not some arbitrary multiple of this - and without reading some documentation its impossible to determine this is what is happening.
MDA simulations have no relation to real world time - the only reason this is an issue is because of the afformetioned PlaybackSimulation.
This is an outlier in the simulation types as it is the only one that has any realation to the real world time elapsing - so I don't see why every other simulation type (that actually implements the method) should bend over backwards to accomodate it.
I would suggest that if this is absolutely required, the onus is on the Runner not the simulation to handle any discrepencies. Or, as I had previously implemented, allow for all 3 options (I'm open to alternative name suggestions).
Further, just because no one else has requested this functionality, doesn't mean its not useful - rather, as I had implemented it, I am implicitly asking to be able to do this.

Regarding handling jumping to a timepoint that is out of bounds for the simulation (e.g. asking to seek_to_time(15) when the simulation is only 10 ns long):
Again, I would argue that the current functionality makes the least sense and is only a product of having the PlaybackSimulation as the previous standard.
I have no issues with this being an option but in most contexts I don't believe this is the intuitive behaviour. I would suggest that jumping to the last timestep (or first) should at least be default (or maybe throw an OutOfBounds style error?).
Consider someone working with the trajectory - without some indicator of time also present (i.e. in VR) - its impossible to tell whether a change you are seeing responds to an actual step of said (nano)seconds. If it simply jumps to a point midway through the simulation (e.g. 5ns in this example) implying the simulation is longer than it is.
Rather, this behaviour should be controlled with some member/setter func to modify this as required.

[Ragzouken]

I have taken the current implementation + tests and merged them (#539), but with the-non protocol methods private so they can easily be renamed at a later date. This puts main in a state where UniverseSimulation supports playback of MDAnalysis universes similarly to recording playback, but without any additional capability in nanover-server command line or exposing non-Simulation features from UniverseSimulation.

What more to add/expose or not e.g advance to next report, seeking, timing control, improving the looping behaviour etc we can continue discussing after Christmas. If I have time I'll try to respond to your points above but our office closes in the next few days.