From c7c251d60a5592e72d351190b961e9288689f5d2 Mon Sep 17 00:00:00 2001
From: richard gowers <richardjgowers@gmail.com>
Date: Thu, 9 Nov 2023 22:39:31 +0000
Subject: [PATCH] cli: add --report raw option to gather

this reports the unit estimates individually
---
 openfecli/commands/gather.py | 33 +++++++++++++++++++++++++++++++--
 1 file changed, 31 insertions(+), 2 deletions(-)

diff --git a/openfecli/commands/gather.py b/openfecli/commands/gather.py
index 1905e971c..d50e25762 100644
--- a/openfecli/commands/gather.py
+++ b/openfecli/commands/gather.py
@@ -128,6 +128,15 @@ def _generate_bad_legs_error_message(set_vals, ligpair):
     return msg
 
 
+def _parse_raw_units(results: dict) -> list[tuple]:
+    # grab individual unit results from master results dict
+    # returns list of (estimate, uncertainty) tuples
+    pus = list(results['unit_results'].values())
+    return [(pu['outputs']['unit_estimate'],
+             pu['outputs']['unit_estimate_error'])
+            for pu in pus]
+
+
 def _get_ddgs(legs, error_on_missing=True):
     import numpy as np
     DDGs = []
@@ -181,6 +190,21 @@ def _write_ddg(legs, writer, allow_partial):
             writer.writerow([ligA, ligB, DDGhyd, hyd_unc])
 
 
+def _write_raw(legs, writer):
+    writer.writerow(["leg", "ligand_i", "ligand_j", "DG(i->j) (kcal/mol)",
+                     "uncertainty (kcal/mol)"])
+
+    for ligpair, vals in sorted(legs.items()):
+        for simtype, repeats in sorted(vals.items()):
+            for m, u in repeats:
+                if m is None:
+                    m, u = 'NaN', 'NaN'
+                else:
+                    m, u = format_estimate_uncertainty(m.m, u.m)
+
+                writer.writerow([simtype, *ligpair, m, u])
+
+
 def _write_dg_raw(legs, writer, allow_partial):
     writer.writerow(["leg", "ligand_i", "ligand_j", "DG(i->j) (kcal/mol)",
                      "uncertainty (kcal/mol)"])
@@ -255,7 +279,8 @@ def _write_dg_mle(legs, writer, allow_partial):
                 required=True)
 @click.option(
     '--report',
-    type=HyphenAwareChoice(['dg', 'ddg', 'dg-raw'], case_sensitive=False),
+    type=HyphenAwareChoice(['dg', 'ddg', 'dg-raw', 'raw'],
+                           case_sensitive=False),
     default="dg", show_default=True,
     help=(
         "What data to report. 'dg' gives maximum-likelihood estimate of "
@@ -325,7 +350,10 @@ def gather(rootdir, output, report, allow_partial):
         except KeyError:
             simtype = legacy_get_type(result_fn)
 
-        legs[names][simtype] = result['estimate'], result['uncertainty']
+        if report.lower() == 'raw':
+            legs[names][simtype] = _parse_raw_units(result)
+        else:
+            legs[names][simtype] = result['estimate'], result['uncertainty']
 
     writer = csv.writer(
         output,
@@ -340,6 +368,7 @@ def gather(rootdir, output, report, allow_partial):
         'dg': _write_dg_mle,
         'ddg': _write_ddg,
         'dg-raw': _write_dg_raw,
+        'raw': _write_raw,
     }[report.lower()]
     writing_func(legs, writer, allow_partial)