generate_free_energy() fails because standard state correction is missing from the reporter

[zhang-ivy]

When running the yank analyze report command one an openmmtools generated repex trajectory, I am getting the following error:

Rendering notebook as a .ipynb file...
Processing notebook now, this may take a while...
Traceback (most recent call last):
  File "/home/zhangi/miniconda3/envs/perses-paper/bin/yank", line 10, in <module>
    sys.exit(main())
  File "/home/zhangi/miniconda3/envs/perses-paper/lib/python3.9/site-packages/yank/cli.py", line 73, in main
    dispatched = getattr(commands, command).dispatch(command_args)
  File "/home/zhangi/miniconda3/envs/perses-paper/lib/python3.9/site-packages/yank/commands/analyze.py", line 123, in dispatch
    return dispatch_report(args)
  File "/home/zhangi/miniconda3/envs/perses-paper/lib/python3.9/site-packages/yank/commands/analyze.py", line 322, in dispatch_report
    run_notebook(store, output, notebook_serial_file, **analyzer_kwargs)
  File "/home/zhangi/miniconda3/envs/perses-paper/lib/python3.9/site-packages/yank/commands/analyze.py", line 281, in run_notebook
    processed_notebook, resources = ep.preprocess(loaded_notebook, resource_data)
  File "/home/zhangi/miniconda3/envs/perses-paper/lib/python3.9/site-packages/nbconvert/preprocessors/execute.py", line 89, in preprocess
    self.preprocess_cell(cell, resources, index)
  File "/home/zhangi/miniconda3/envs/perses-paper/lib/python3.9/site-packages/nbconvert/preprocessors/execute.py", line 110, in preprocess_cell
    cell = self.execute_cell(cell, index, store_history=True)
  File "/home/zhangi/miniconda3/envs/perses-paper/lib/python3.9/site-packages/nbclient/util.py", line 85, in wrapped
    return just_run(coro(*args, **kwargs))
  File "/home/zhangi/miniconda3/envs/perses-paper/lib/python3.9/site-packages/nbclient/util.py", line 60, in just_run
    return loop.run_until_complete(coro)
  File "/home/zhangi/miniconda3/envs/perses-paper/lib/python3.9/asyncio/base_events.py", line 647, in run_until_complete
    return future.result()
  File "/home/zhangi/miniconda3/envs/perses-paper/lib/python3.9/site-packages/nbclient/client.py", line 1019, in async_execute_cell
    await self._check_raise_for_error(cell, cell_index, exec_reply)
  File "/home/zhangi/miniconda3/envs/perses-paper/lib/python3.9/site-packages/nbclient/client.py", line 913, in _check_raise_for_error
    raise CellExecutionError.from_cell_and_msg(cell, exec_reply_content)
nbclient.exceptions.CellExecutionError: An error occurred while executing the following cell:
------------------
report.generate_free_energy()
------------------

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
Input In [10], in <cell line: 1>()
----> 1 report.generate_free_energy()

File ~/miniconda3/envs/perses-paper/lib/python3.9/site-packages/yank/reports/notebook.py:491, in HealthReportData.generate_free_energy(self)
    490 def generate_free_energy(self):
--> 491     fe_data = self.get_experiment_free_energy_data()
    492     delta_f = fe_data['free_energy_diff']
    493     delta_h = fe_data['enthalpy_diff']

File ~/miniconda3/envs/perses-paper/lib/python3.9/site-packages/yank/analyze.py:177, in _copyout.<locals>.make_copy(*args, **kwargs)
    175 @wraps(wrap)
    176 def make_copy(*args, **kwargs):
--> 177     return copy.deepcopy(wrap(*args, **kwargs))

File ~/miniconda3/envs/perses-paper/lib/python3.9/site-packages/yank/analyze.py:569, in ExperimentAnalyzer.get_experiment_free_energy_data(self)
    567 for phase_name in self.phase_names:
    568     analyzer = self.analyzers[phase_name]
--> 569     data[phase_name] = analyzer.analyze_phase()
    571 # Compute free energy and enthalpy
    572 delta_f = 0.0

File ~/miniconda3/envs/perses-paper/lib/python3.9/site-packages/yank/analyze.py:161, in YankMultiStateSamplerAnalyzer.analyze_phase(self, show_mixing, cutoff)
    159 data['enthalpy_diff'] = DeltaH_ij[self.reference_states[0], self.reference_states[1]]
    160 data['enthalpy_diff_error'] = dDeltaH_ij[self.reference_states[0], self.reference_states[1]]
--> 161 data['free_energy_diff_standard_state_correction'] = self.get_standard_state_correction()
    162 return data

File ~/miniconda3/envs/perses-paper/lib/python3.9/site-packages/yank/analyze.py:114, in YankMultiStateSamplerAnalyzer.get_standard_state_correction(self)
    111 # Reads the SSC from the reporter if compute_ssc is False.
    112 if self._computed_observables['standard_state_correction'] is None:
    113     #try:
--> 114     ssc = self._reporter.read_dict('metadata')['standard_state_correction']
    115     #except:
    116     #    logger.debug("standard state correction was not present in the reporter metadata, setting it to 0.0")
    117     #    ssc = 0.0
    118     self._computed_observables['standard_state_correction'] = ssc

KeyError: 'standard_state_correction'
KeyError: 'standard_state_correction'

The error occurs because of YankMultiStateSamplerAnalyzer.get_standard_state_correction() (linked here). compute_ssc is set to False here because there is no restraint force and then it attempts to read the standard state correction from the reporter here, but that's not a variable that's saved when running simulations in openmmtools.

This is the fix I added to my local version of yank/analyze.py:

# Reads the SSC from the reporter if compute_ssc is False.
if self._computed_observables['standard_state_correction'] is None:
    try:
        ssc = self._reporter.read_dict('metadata')['standard_state_correction']
    except:
        logger.debug("standard state correction was not present in the reporter metadata, setting it to 0.0")
        ssc = 0.0
    self._computed_observables['standard_state_correction'] = ssc

I could open a PR with this fix, but I'm not sure if we are actively maintaining yank / cutting a release soon?

@jchodera mentioned awhile ago that eventually we'd want to move all the yank analysis stuff to openmmtools, but I'm wondering what the best short term solution here is.

cc: @mikemhenry @ijpulidos for feedback on this too

[mikemhenry]

That looks reasonable, I would say that you are better off with except KeyError so that you only catch that exception instead of maybe something else that could be thrown. "Naked" excepts can make code really hard to debug later.