ADicksonLab / wepy

Weighted Ensemble simulation framework in Python
https://adicksonlab.github.io/wepy/index.html
MIT License
51 stars 20 forks source link

Refactor and consistent handling of OpenMMState fields #117

Open salotz opened 11 months ago

salotz commented 11 months ago

Currently the mechanisms for wrapping and retrieving dynamically available fields in the OpenMMState and how it wraps the OpenMM.State is very flaky and causes problems for the tail of attributes that aren't actually ever used in practice very much (e.g. parameter derivatives).

I would like to refactor this so that these only cause problems when a user actually wants these things.

A checklist:

I have some work towards this already and will need to just finish some things on this checklist to maintain full support and backwards compatibility.

alexrd commented 10 months ago

Adding some urgency to this, as it seems that the latest openmm will raise an exception when we call getEnergyParameterDerivatives on a state that doesn't have them:

Traceback (most recent call last):
...
  File "/mnt/ufs18/home-126/alexrd/wepy/src/wepy/sim_manager.py", line 492, in _run_cycle
    reporter.report(**report)
  File "/mnt/ufs18/home-126/alexrd/wepy/src/wepy/reporter/hdf5.py", line 523, in report
    walker_data = walker.state.dict()
  File "/mnt/ufs18/home-126/alexrd/wepy/src/wepy/runners/openmm.py", line 1194, in dict
    for key, value in self.omm_state_dict().items():
  File "/mnt/ufs18/home-126/alexrd/wepy/src/wepy/runners/openmm.py", line 1179, in omm_state_dict
    param_derivs = self.parameter_derivatives_features()
  File "/mnt/ufs18/home-126/alexrd/wepy/src/wepy/runners/openmm.py", line 1155, in parameter_derivatives_features
    parameter_derivatives = self.parameter_derivatives_values()
  File "/mnt/ufs18/home-126/alexrd/wepy/src/wepy/runners/openmm.py", line 1047, in parameter_derivatives_values
    if self.parameter_derivatives is None:
  File "/mnt/ufs18/home-126/alexrd/wepy/src/wepy/runners/openmm.py", line 1020, in parameter_derivatives
    return self.sim_state.getEnergyParameterDerivatives()
  File "/mnt/home/alexrd/anaconda3/envs/wepy/lib/python3.10/site-packages/openmm/openmm.py", line 4965, in getEnergyParameterDerivatives
    return _openmm.State_getEnergyParameterDerivatives(self)
openmm.OpenMMException: Invoked getEnergyParameterDerivatives() on a State which does not contain parameter derivatives.
salotz commented 10 months ago

Yes this is the error I was finding as well. Interestingly it doesn't show up when using Reference or CPU platform.

I'll try to merge a comprehensive fix of this very soon.