Use specialized reporter for production trajectories

[mattwthompson]

In our notes I have

Trajectory reporting - need to use Simon’s specialized trajectory class in (it’s something msgpack in descent)

I'm pretty sure this is referring to smee.mm.TensorReporter

But for now keep both this and DCD files around. (Long-term maybe also turn on DCD files for debugging)

[mattwthompson]

Below is based on #113 at 7d71b3f, reproduce using pixi run -e dev ipython examples/run_density.py.

I don't know if the failure is openmm.unit.Quantity or openff.toolkit.Quantity but only OpenMM quantities are pass to the reporter so I figure it can only be that.

It's possible the failure is related to Parsl's pickling of intermediate objects, but my reading of the traceback is that it's just in OpenMM world.

DEBUG:parsl.app.errors:Reraising exception of type <class 'TypeError'>
INFO:parsl.dataflow.dflow:Task 3 try 0 failed with exception of type TypeError
ERROR:parsl.dataflow.dflow:Task 3 failed after 0 retry attempts
Traceback (most recent call last):
  File "/Users/mattthompson/software/dimsim/.pixi/envs/dev/lib/python3.12/site-packages/parsl/dataflow/dflow.py", line 348, in handle_exec_update
    res = self._unwrap_remote_exception_wrapper(future)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/mattthompson/software/dimsim/.pixi/envs/dev/lib/python3.12/site-packages/parsl/dataflow/dflow.py", line 578, in _unwrap_remote_exception_wrapper
    result.reraise()
  File "/Users/mattthompson/software/dimsim/.pixi/envs/dev/lib/python3.12/site-packages/parsl/app/errors.py", line 114, in reraise
    raise v
  File "/Users/mattthompson/software/dimsim/.pixi/envs/dev/lib/python3.12/site-packages/parsl/app/errors.py", line 138, in wrapper
    return func(*args, **kwargs)
    ^^^^^^^^^^^^^^^^^
  File "/Users/mattthompson/software/dimsim/dimsim/compute/apps.py", line 286, in run_equilibration
    simulation.step(10_000)
      ^^^^^^^^^^^^^^^^^
  File "/Users/mattthompson/software/dimsim/.pixi/envs/dev/lib/python3.12/site-packages/openmm/app/simulation.py", line 149, in step
    self._simulate(endStep=self.currentStep+steps)
    ^^^^^^^^^^^^^^^^^
  File "/Users/mattthompson/software/dimsim/.pixi/envs/dev/lib/python3.12/site-packages/openmm/app/simulation.py", line 260, in _simulate
    self._generate_reports(wrapped, True)
    ^^^^^^^^^^^^^^^^^
  File "/Users/mattthompson/software/dimsim/.pixi/envs/dev/lib/python3.12/site-packages/openmm/app/simulation.py", line 281, in _generate_reports
    reporter.report(self, state)
    ^^^^^^^^^^^^^^^^^
  File "/Users/mattthompson/software/dimsim/.pixi/envs/dev/lib/python3.12/site-packages/smee/mm/_reporters.py", line 104, in report
    self._output_file.write(msgpack.dumps(frame, default=_encoder))
    ^^^^^^^^^^^^^^^^^
  File "/Users/mattthompson/software/dimsim/.pixi/envs/dev/lib/python3.12/site-packages/msgpack/__init__.py", line 36, in packb
    return Packer(**kwargs).pack(o)
      ^^^^^^^^^^^^^^^^^
  File "msgpack/_packer.pyx", line 279, in msgpack._cmsgpack.Packer.pack
  File "msgpack/_packer.pyx", line 276, in msgpack._cmsgpack.Packer.pack
  File "msgpack/_packer.pyx", line 264, in msgpack._cmsgpack.Packer._pack
  File "msgpack/_packer.pyx", line 231, in msgpack._cmsgpack.Packer._pack_inner
  File "msgpack/_packer.pyx", line 269, in msgpack._cmsgpack.Packer._pack
  File "msgpack/_packer.pyx", line 256, in msgpack._cmsgpack.Packer._pack_inner
TypeError: can not serialize 'Quantity' object

[lilyminium]

Beta should be in units of 1/energy (i.e. 1/kilocalories_per_mole), could that be an issue? I think OpenMM units indeed can't be serialized, I'm guessing the reduced_potential should get reduced to dimensionless (and perhaps implicitly converted to floats) but are staying as quantities here.

Otherwise smee might need an update to only save reduced_potential.value_in_unit(dimensionless) to work here.

[mattwthompson]

Thanks @lilyminium! I think that was it. With that simple fix my example script runs through again.

I think that the error is a downstream effect of the quirk that OpenMM turns dimensionless quantities into raw floats, etc. when their units cancel each other out. So when that doesn't happen, it runs into the task of serializing the quantity object which it hadn't needed to do before.

This is all assuming that the reporter works fine out of the box. I haven't really used it yet so I can't rule out other improvements that might need to be made