Add gather --allow-partial, improve related error message (#588)

* improve error message

* Add --allow-partial to gather

* use it for the other call as well

* ignore coverage on line that should never happen

* Make it possible to do both RHFE and RBFE

openfecli/commands/gather.py

@@ -5,6 +5,7 @@
 from openfecli import OFECommandPlugin
 from openfecli.clicktypes import HyphenAwareChoice
 import pathlib
+import warnings
 
 
 def _get_column(val):
@@ -85,40 +86,90 @@ def legacy_get_type(res_fn):
  return 'complex'
 
 
-def _get_ddgs(legs):
+def _generate_bad_legs_error_message(set_vals, ligpair):
+ expected_rbfe = {'complex', 'solvent'}
+ expected_rhfe = {'solvent', 'vacuum'}
+ maybe_rhfe = bool(set_vals & expected_rhfe)
+ maybe_rbfe = bool(set_vals & expected_rbfe)
+ if maybe_rhfe and not maybe_rbfe:
+ msg = (
+ "This appears to be an RHFE calculation, but we're "
+ f"missing {expected_rhfe - set_vals} runs for the "
+ f"edge with ligands {ligpair}."
+ )
+ elif maybe_rbfe and not maybe_rhfe:
+ msg = (
+ "This appears to be an RBFE calculation, but we're "
+ f"missing {expected_rbfe - set_vals} runs for the "
+ f"edge with ligands {ligpair}."
+ )
+ elif maybe_rbfe and maybe_rhfe:
+ msg = (
+ "Unable to determine whether this is an RBFE "
+ f"or an RHFE calculation. Found legs {set_vals} "
+ f"for ligands {ligpair}. Those ligands are missing one "
+ f"of: {(expected_rhfe | expected_rbfe) - set_vals}."
+ )
+ else: # -no-cov-
+ # this should never happen
+ msg = (
+ "Something went very wrong while determining the type "
+ f"of RFE calculation. For the ligand pair {ligpair}, "
+ f"we found legs labelled {set_vals}. We expected either "
+ f"{expected_rhfe} or {expected_rbfe}."
+ )
+
+ msg += (
+ "\n\nYou can force partial gathering of results, without "
+ "problematic edges, by using the --allow-partial flag of the gather "
+ "command. Note that this may cause problems with predicting "
+ "absolute free energies from the relative free energies."
+ )
+ return msg
+
+
+def _get_ddgs(legs, error_on_missing=True):
  import numpy as np
  DDGs = []
  for ligpair, vals in sorted(legs.items()):
+ set_vals = set(vals)
  DDGbind = None
  DDGhyd = None
  bind_unc = None
  hyd_unc = None
 
- if 'complex' in vals and 'solvent' in vals:
+ do_rbfe = (len(set_vals & {'complex', 'solvent'}) == 2)
+ do_rhfe = (len(set_vals & {'vacuum', 'solvent'}) == 2)
+
+ if do_rbfe:
  DG1_mag, DG1_unc = vals['complex']
  DG2_mag, DG2_unc = vals['solvent']
  if not ((DG1_mag is None) or (DG2_mag is None)):
  # DDG(2,1)bind = DG(1->2)complex - DG(1->2)solvent
  DDGbind = (DG1_mag - DG2_mag).m
  bind_unc = np.sqrt(np.sum(np.square([DG1_unc.m, DG2_unc.m])))
- elif 'solvent' in vals and 'vacuum' in vals:
+
+ if do_rhfe:
  DG1_mag, DG1_unc = vals['solvent']
  DG2_mag, DG2_unc = vals['vacuum']
  if not ((DG1_mag is None) or (DG2_mag is None)):
  DDGhyd = (DG1_mag - DG2_mag).m
  hyd_unc = np.sqrt(np.sum(np.square([DG1_unc.m, DG2_unc.m])))
- else:
- raise RuntimeError("Unable to determine type of RFE calculation "
- f"for edges with labels {list(vals)} for "
- f"ligands {ligpair}")
+
+ if not do_rbfe and not do_rhfe:
+ msg = _generate_bad_legs_error_message(set_vals, ligpair)
+ if error_on_missing:
+ raise RuntimeError(msg)
+ else:
+ warnings.warn(msg)
 
  DDGs.append((*ligpair, DDGbind, bind_unc, DDGhyd, hyd_unc))
 
  return DDGs
 
 
-def _write_ddg(legs, writer):
- DDGs = _get_ddgs(legs)
+def _write_ddg(legs, writer, allow_partial):
+ DDGs = _get_ddgs(legs, error_on_missing=not allow_partial)
  writer.writerow(["ligand_i", "ligand_j", "DDG(i->j) (kcal/mol)",
  "uncertainty (kcal/mol)"])
  for ligA, ligB, DDGbind, bind_unc, DDGhyd, hyd_unc in DDGs:
@@ -133,7 +184,7 @@ def _write_ddg(legs, writer):
  writer.writerow([ligA, ligB, DDGhyd, hyd_unc])
 
 
-def _write_dg_raw(legs, writer):
+def _write_dg_raw(legs, writer, allow_partial):
  writer.writerow(["leg", "ligand_i", "ligand_j", "DG(i->j) (kcal/mol)",
  "uncertainty (kcal/mol)"])
  for ligpair, vals in sorted(legs.items()):
@@ -146,11 +197,11 @@ def _write_dg_raw(legs, writer):
  writer.writerow([simtype, *ligpair, m, u])
 
 
-def _write_dg_mle(legs, writer):
+def _write_dg_mle(legs, writer, allow_partial):
  import networkx as nx
  import numpy as np
  from cinnabar.stats import mle
- DDGs = _get_ddgs(legs)
+ DDGs = _get_ddgs(legs, error_on_missing=not allow_partial)
  MLEs = []
  # 4b) perform MLE
  g = nx.DiGraph()
@@ -219,7 +270,14 @@ def _write_dg_mle(legs, writer):
 @click.option('output', '-o',
  type=click.File(mode='w'),
  default='-')
-def gather(rootdir, output, report):
+@click.option(
+ '--allow-partial', is_flag=True, default=False,
+ help=(
+ "Do not raise errors is results are missing parts for some edges. "
+ "(Skip those edges and issue warning instead.)"
+ )
+)
+def gather(rootdir, output, report, allow_partial):
  """Gather simulation result jsons of relative calculations to a tsv file
 
  This walks ROOTDIR recursively and finds all result JSON files from the
@@ -287,7 +345,7 @@ def gather(rootdir, output, report):
  'ddg': _write_ddg,
  'dg-raw': _write_dg_raw,
  }[report.lower()]
- writing_func(legs, writer)
+ writing_func(legs, writer, allow_partial)
 
 
 PLUGIN = OFECommandPlugin(

openfecli/tests/commands/test_gather.py

@@ -6,7 +6,8 @@
 import pytest
 
 from openfecli.commands.gather import (
- gather, format_estimate_uncertainty, _get_column
+ gather, format_estimate_uncertainty, _get_column,
+ _generate_bad_legs_error_message,
 )
 
 @pytest.mark.parametrize('est,unc,unc_prec,est_str,unc_str', [
@@ -108,6 +109,21 @@ def test_gather(results_dir, report):
  assert set(expected.split(b'\n')) == actual_lines
 
 
+@pytest.mark.parametrize('include', ['complex', 'solvent', 'vacuum'])
+def test_generate_bad_legs_error_message(include):
+ expected = {
+ 'complex': ("appears to be an RBFE", "missing {'solvent'}"),
+ 'vacuum': ("appears to be an RHFE", "missing {'solvent'}"),
+ 'solvent': ("whether this is an RBFE or an RHFE",
+ "'complex'", "'solvent'"),
+ }[include]
+ set_vals = {include}
+ ligpair = {'lig1', 'lig2'}
+ msg = _generate_bad_legs_error_message(set_vals, ligpair)
+ for string in expected:
+ assert string in msg
+
+
 def test_missing_leg_error(results_dir):
  file_to_remove = "easy_rbfe_lig_ejm_31_complex_lig_ejm_42_complex.json"
  (pathlib.Path("results") / file_to_remove).unlink()
@@ -116,6 +132,16 @@ def test_missing_leg_error(results_dir):
  result = runner.invoke(gather, ['results'] + ['-o', '-'])
  assert result.exit_code == 1
  assert isinstance(result.exception, RuntimeError)
- assert "labels ['solvent']" in str(result.exception)
+ assert "Unable to determine" in str(result.exception)
  assert "'lig_ejm_31'" in str(result.exception)
  assert "'lig_ejm_42'" in str(result.exception)
+
+
+def test_missing_leg_allow_partial(results_dir):
+ file_to_remove = "easy_rbfe_lig_ejm_31_complex_lig_ejm_42_complex.json"
+ (pathlib.Path("results") / file_to_remove).unlink()
+
+ runner = CliRunner()
+ result = runner.invoke(gather,
+ ['results'] + ['--allow-partial', '-o', '-'])
+ assert result.exit_code == 0