From aa73c0c6e88013a258c7dc412fb1c8ce17210827 Mon Sep 17 00:00:00 2001 From: Alyssa Travitz Date: Wed, 14 Jan 2026 08:05:43 -0800 Subject: [PATCH 1/8] refactor quickrun data into dataclass --- openfecli/commands/quickrun.py | 60 +++++++++++++++++++++++++++------- 1 file changed, 49 insertions(+), 11 deletions(-) diff --git a/openfecli/commands/quickrun.py b/openfecli/commands/quickrun.py index f34410d69..baf8674de 100644 --- a/openfecli/commands/quickrun.py +++ b/openfecli/commands/quickrun.py @@ -3,8 +3,12 @@ import json import pathlib +from dataclasses import asdict, dataclass +from typing import Any, Self import click +from gufe.tokenization import JSON_HANDLER +from openff.units import Quantity from openfecli import OFECommandPlugin from openfecli.utils import configure_logger, print_duration, write @@ -15,6 +19,42 @@ def _format_exception(exception) -> str: return f"{exception[0]}: {exception[1][0]}" +@dataclass +class QuickrunResult: + """ + Class for storing protocol result data along with useful metadata. + Could ProtocolResults store this data alongside ``n_protocol_dag_results``? + """ + + estimate: Quantity + uncertainty: Quantity + protocol_result: dict[str, Any] + unit_results: dict[int, dict] + # ok: bool # TODO: add in 2.0? + + def to_json(self, filepath) -> None: + with open(filepath, mode="w") as file: + json.dump(asdict(self), file, cls=JSON_HANDLER.encoder) + + @classmethod + def from_json(cls, filepath) -> Self: + """Load a JSON file containing a gufe object. + + Parameters + ---------- + fpath : os.PathLike | str + The path to a results JSON generated by ``openfe quickrun``. + + + Returns + ------- + QuickrunResult + A QuickrunResult instance containing the data from ``filepath``. + + """ + return json.load(open(filepath, "r"), cls=JSON_HANDLER.decoder) + + @click.command("quickrun", short_help="Run a given transformation, saved as a JSON file") @click.argument("transformation", type=click.File(mode="r"), required=True) @click.option( @@ -114,17 +154,15 @@ def quickrun(transformation, work_dir, output): else: estimate = uncertainty = None # for output file - out_dict = { - "estimate": estimate, - "uncertainty": uncertainty, - "protocol_result": prot_result.to_dict(), - "unit_results": { - unit.key: unit.to_keyed_dict() for unit in dagresult.protocol_unit_results - }, - } - - with open(output, mode="w") as outf: - json.dump(out_dict, outf, cls=JSON_HANDLER.encoder) + quickrun_result = QuickrunResult( + estimate=estimate, + uncertainty=uncertainty, + protocol_result=prot_result.to_dict(), + unit_results={unit.key: unit.to_keyed_dict() for unit in dagresult.protocol_unit_results}, + # ok=dagresult.ok() # TODO: add in 2.0 + ) + + quickrun_result.to_json(output) write(f"Here is the result:\n\tdG = {estimate} ± {uncertainty}\n") write("") From 0c152cc88e9a2740fd8973d75ad4680c05137767 Mon Sep 17 00:00:00 2001 From: Alyssa Travitz Date: Wed, 14 Jan 2026 09:27:48 -0800 Subject: [PATCH 2/8] update tests to use dataclass --- openfecli/commands/gather.py | 29 ++++++----- openfecli/commands/quickrun.py | 16 ++++-- openfecli/tests/commands/test_gather.py | 69 +++++++++++++++++-------- 3 files changed, 76 insertions(+), 38 deletions(-) diff --git a/openfecli/commands/gather.py b/openfecli/commands/gather.py index a99e4ee8e..0226096b0 100644 --- a/openfecli/commands/gather.py +++ b/openfecli/commands/gather.py @@ -11,6 +11,7 @@ from openfecli import OFECommandPlugin from openfecli.clicktypes import HyphenAwareChoice +from openfecli.commands.quickrun import QuickrunResult FAIL_STR = "Error" # string used to indicate a failed run in output tables. @@ -204,7 +205,7 @@ def load_json(fpath: os.PathLike | str) -> dict: return json.load(open(fpath, "r"), cls=JSON_HANDLER.decoder) -def _get_names(result: dict) -> tuple[str, str]: +def _get_names(result: QuickrunResult) -> tuple[str, str]: """Get the ligand names from a unit's results data. Parameters @@ -219,7 +220,7 @@ def _get_names(result: dict) -> tuple[str, str]: """ # TODO: I don't like this [0][0] indexing, but I can't think of a better way currently - protocol_data = list(result["protocol_result"]["data"].values())[0][0] + protocol_data = list(result.protocol_result["data"].values())[0][0] name_A = protocol_data["inputs"]["ligandmapping"]["componentA"]["molprops"]["ofe-name"] name_B = protocol_data["inputs"]["ligandmapping"]["componentB"]["molprops"]["ofe-name"] @@ -227,10 +228,10 @@ def _get_names(result: dict) -> tuple[str, str]: return str(name_A), str(name_B) -def _get_type(result: dict) -> Literal["vacuum", "solvent", "complex"]: +def _get_type(result: QuickrunResult) -> Literal["vacuum", "solvent", "complex"]: """Determine the simulation type based on the component types.""" - protocol_data = list(result["protocol_result"]["data"].values())[0][0] + protocol_data = list(result.protocol_result["data"].values())[0][0] component_types = [ x["__module__"] for x in protocol_data["inputs"]["stateA"]["components"].values() ] @@ -255,7 +256,7 @@ def _legacy_get_type(res_fn: os.PathLike | str) -> Literal["vacuum", "solvent", def _get_result_id( - result: dict, result_fn: os.PathLike | str + result: QuickrunResult, result_fn: os.PathLike | str ) -> tuple[tuple[str, str], Literal["vacuum", "solvent", "complex"]]: """Extract the name and simulation type from a results dict. @@ -281,7 +282,7 @@ def _get_result_id( return (ligA, ligB), simtype -def _load_valid_result_json(fpath: os.PathLike | str) -> tuple[tuple | None, dict | None]: +def _load_valid_result_json(fpath: os.PathLike | str) -> tuple[tuple | None, QuickrunResult | None]: """Load the data from a results JSON into a dict. Parameters @@ -296,25 +297,25 @@ def _load_valid_result_json(fpath: os.PathLike | str) -> tuple[tuple | None, dic or None if the JSON file is invalid or missing. """ - # TODO: only load this once during collection, then pass namedtuple(fname, dict) into this function # for now though, it's not the bottleneck on performance - result = load_json(fpath) + result = QuickrunResult.from_json(fpath) + try: result_id = _get_result_id(result, fpath) except (ValueError, IndexError): click.secho(f"{fpath}: Missing ligand names and/or simulation type. Skipping.",err=True, fg="yellow") # fmt: skip return None, None - if result["estimate"] is None: + if result.estimate is None: click.secho(f"{fpath}: No 'estimate' found, assuming to be a failed simulation.",err=True, fg="yellow") # fmt: skip return result_id, None - if result["uncertainty"] is None: + if result.uncertainty is None: click.secho(f"{fpath}: No 'uncertainty' found, assuming to be a failed simulation.",err=True, fg="yellow") # fmt: skip return result_id, None - if result["unit_results"] == {}: + if result.unit_results == {}: click.secho(f"{fpath}: No 'unit_results' found, assuming to be a failed simulation.",err=True, fg="yellow") # fmt: skip return result_id, None - if all("exception" in u for u in result["unit_results"].values()): + if all("exception" in u for u in result.unit_results.values()): click.secho(f"{fpath}: Exception found in all 'unit_results', assuming to be a failed simulation.",err=True, fg="yellow") # fmt: skip return result_id, None @@ -667,7 +668,7 @@ def _get_legs_from_result_jsons( v[0]["outputs"]["unit_estimate"], v[0]["outputs"]["unit_estimate_error"], ) - for v in result["protocol_result"]["data"].values() + for v in result.protocol_result["data"].values() ] legs[names][simtype].append(parsed_raw_data) else: @@ -677,7 +678,7 @@ def _get_legs_from_result_jsons( else: dGs = [ v[0]["outputs"]["unit_estimate"] - for v in result["protocol_result"]["data"].values() + for v in result.protocol_result["data"].values() ] legs[names][simtype].extend(dGs) return legs diff --git a/openfecli/commands/quickrun.py b/openfecli/commands/quickrun.py index baf8674de..9986b118e 100644 --- a/openfecli/commands/quickrun.py +++ b/openfecli/commands/quickrun.py @@ -4,6 +4,7 @@ import json import pathlib from dataclasses import asdict, dataclass +from os import PathLike from typing import Any, Self import click @@ -37,7 +38,7 @@ def to_json(self, filepath) -> None: json.dump(asdict(self), file, cls=JSON_HANDLER.encoder) @classmethod - def from_json(cls, filepath) -> Self: + def from_json(cls, file: PathLike | None, content: str | None = None) -> Self: """Load a JSON file containing a gufe object. Parameters @@ -49,10 +50,19 @@ def from_json(cls, filepath) -> Self: Returns ------- QuickrunResult - A QuickrunResult instance containing the data from ``filepath``. + A QuickrunResult instance containing the data from ``file, co``. """ - return json.load(open(filepath, "r"), cls=JSON_HANDLER.decoder) + # similar to gufe.tokenization.from_json + if content is not None and file is not None: + raise ValueError("Cannot specify both `content` and `file`; only one input allowed") + elif content is None and file is None: + raise ValueError("Must specify either `content` and `file` for JSON input") + if file: + data = json.load(open(file, "r"), cls=JSON_HANDLER.decoder) + if content: + data = json.loads(content, cls=JSON_HANDLER.decoder) + return cls(**data) @click.command("quickrun", short_help="Run a given transformation, saved as a JSON file") diff --git a/openfecli/tests/commands/test_gather.py b/openfecli/tests/commands/test_gather.py index 9ea27b59b..7aa149b49 100644 --- a/openfecli/tests/commands/test_gather.py +++ b/openfecli/tests/commands/test_gather.py @@ -17,6 +17,7 @@ ) from openfecli.commands.gather_abfe import gather_abfe from openfecli.commands.gather_septop import gather_septop +from openfecli.commands.quickrun import QuickrunResult from ..conftest import HAS_INTERNET from ..utils import assert_click_success @@ -70,62 +71,88 @@ def test_get_column(val, col): assert _get_column(val) == col +@pytest.fixture +def min_valid_quickrun_result(min_result_json): + return QuickrunResult(**min_result_json) + + class TestResultLoading: - def test_minimal_valid_results(self, capsys, min_result_json): - with mock.patch("openfecli.commands.gather.load_json", return_value=min_result_json): + def test_minimal_valid_results(self, capsys, min_valid_quickrun_result): + with mock.patch( + "openfecli.commands.gather.QuickrunResult.from_json", + return_value=min_valid_quickrun_result, + ): result = _load_valid_result_json(fpath="") captured = capsys.readouterr() - assert result == ((("lig_ejm_31", "lig_ejm_42"), "solvent"), min_result_json) + assert result == ((("lig_ejm_31", "lig_ejm_42"), "solvent"), min_valid_quickrun_result) assert captured.err == "" - def test_skip_missing_unit_result(self, capsys, min_result_json): - min_result_json["unit_results"] = {} + def test_skip_missing_unit_result(self, capsys, min_valid_quickrun_result): + min_valid_quickrun_result.unit_results = {} - with mock.patch("openfecli.commands.gather.load_json", return_value=min_result_json): + with mock.patch( + "openfecli.commands.gather.QuickrunResult.from_json", + return_value=min_valid_quickrun_result, + ): result = _load_valid_result_json(fpath="") captured = capsys.readouterr() assert result == ((("lig_ejm_31", "lig_ejm_42"), "solvent"), None) assert "No 'unit_results' found" in captured.err - def test_skip_missing_estimate(self, capsys, min_result_json): - min_result_json["estimate"] = None + def test_skip_missing_estimate(self, capsys, min_valid_quickrun_result): + min_valid_quickrun_result.estimate = None - with mock.patch("openfecli.commands.gather.load_json", return_value=min_result_json): + with mock.patch( + "openfecli.commands.gather.QuickrunResult.from_json", + return_value=min_valid_quickrun_result, + ): result = _load_valid_result_json(fpath="") captured = capsys.readouterr() assert result == ((("lig_ejm_31", "lig_ejm_42"), "solvent"), None) assert "No 'estimate' found" in captured.err - def test_skip_missing_uncertainty(self, capsys, min_result_json): - min_result_json["uncertainty"] = None + def test_skip_missing_uncertainty(self, capsys, min_valid_quickrun_result): + min_valid_quickrun_result.uncertainty = None - with mock.patch("openfecli.commands.gather.load_json", return_value=min_result_json): + with mock.patch( + "openfecli.commands.gather.QuickrunResult.from_json", + return_value=min_valid_quickrun_result, + ): result = _load_valid_result_json(fpath="") captured = capsys.readouterr() assert result == ((("lig_ejm_31", "lig_ejm_42"), "solvent"), None) assert "No 'uncertainty' found" in captured.err - def test_skip_all_failed_runs(self, capsys, min_result_json): - del min_result_json["unit_results"]["ProtocolUnitResult-e85"] - with mock.patch("openfecli.commands.gather.load_json", return_value=min_result_json): + def test_skip_all_failed_runs(self, capsys, min_valid_quickrun_result): + del min_valid_quickrun_result.unit_results["ProtocolUnitResult-e85"] + with mock.patch( + "openfecli.commands.gather.QuickrunResult.from_json", + return_value=min_valid_quickrun_result, + ): result = _load_valid_result_json(fpath="") captured = capsys.readouterr() assert result == ((("lig_ejm_31", "lig_ejm_42"), "solvent"), None) assert "Exception found in all" in captured.err - def test_missing_pr_data(self, capsys, min_result_json): - min_result_json["protocol_result"]["data"] = {} - with mock.patch("openfecli.commands.gather.load_json", return_value=min_result_json): + def test_missing_pr_data(self, capsys, min_valid_quickrun_result): + min_valid_quickrun_result.protocol_result["data"] = {} + with mock.patch( + "openfecli.commands.gather.QuickrunResult.from_json", + return_value=min_valid_quickrun_result, + ): result = _load_valid_result_json(fpath="") captured = capsys.readouterr() assert result == (None, None) assert "Missing ligand names and/or simulation type. Skipping" in captured.err - def test_get_legs_from_result_jsons(self, capsys, min_result_json): + def test_get_legs_from_result_jsons(self, capsys, min_valid_quickrun_result): """Test that exceptions are handled correctly at the _get_legs_from_results_json level.""" - min_result_json["protocol_result"]["data"] = {} + min_valid_quickrun_result.protocol_result["data"] = {} - with mock.patch("openfecli.commands.gather.load_json", return_value=min_result_json): + with mock.patch( + "openfecli.commands.gather.QuickrunResult.from_json", + return_value=min_valid_quickrun_result, + ): result = _get_legs_from_result_jsons(result_fns=[""], report="dg") captured = capsys.readouterr() assert result == {} From 852004e2744753f4ab7aef513be8fe9899befb40 Mon Sep 17 00:00:00 2001 From: Alyssa Travitz Date: Wed, 14 Jan 2026 09:59:33 -0800 Subject: [PATCH 3/8] update tests to use dataclass --- openfecli/commands/gather.py | 23 --------------------- openfecli/commands/gather_abfe.py | 32 ++++++++++++++--------------- openfecli/commands/gather_septop.py | 32 ++++++++++++++--------------- 3 files changed, 32 insertions(+), 55 deletions(-) diff --git a/openfecli/commands/gather.py b/openfecli/commands/gather.py index 0226096b0..7bc296899 100644 --- a/openfecli/commands/gather.py +++ b/openfecli/commands/gather.py @@ -182,29 +182,6 @@ def is_results_json(fpath: os.PathLike | str) -> bool: return "estimate" in open(fpath, "r").read(20) -def load_json(fpath: os.PathLike | str) -> dict: - """Load a JSON file containing a gufe object. - - Parameters - ---------- - fpath : os.PathLike | str - The path to a gufe-serialized JSON. - - - Returns - ------- - dict - A dict containing data from the results JSON. - - """ - # TODO: move this function to openfe/utils - import json - - from gufe.tokenization import JSON_HANDLER - - return json.load(open(fpath, "r"), cls=JSON_HANDLER.decoder) - - def _get_names(result: QuickrunResult) -> tuple[str, str]: """Get the ligand names from a unit's results data. diff --git a/openfecli/commands/gather_abfe.py b/openfecli/commands/gather_abfe.py index a541bfc52..26a02de9e 100644 --- a/openfecli/commands/gather_abfe.py +++ b/openfecli/commands/gather_abfe.py @@ -13,12 +13,12 @@ from openfecli.commands.gather import ( _collect_result_jsons, format_df_with_precision, - load_json, rich_print_to_stdout, ) +from openfecli.commands.quickrun import QuickrunResult -def _get_name(result: dict) -> str: +def _get_name(result: QuickrunResult) -> str: """Get the ligand name from a unit's results data. Parameters @@ -32,13 +32,13 @@ def _get_name(result: dict) -> str: Ligand name corresponding to the results. """ - solvent_data = list(result["protocol_result"]["data"]["solvent"].values())[0][0] + solvent_data = list(result.protocol_result["data"]["solvent"].values())[0][0] name = solvent_data["inputs"]["alchemical_components"]["stateA"][0]["molprops"]["ofe-name"] return str(name) -def _load_valid_result_json(fpath: os.PathLike | str) -> tuple[tuple | None, dict | None]: +def _load_valid_result_json(fpath: os.PathLike | str) -> tuple[tuple | None, QuickrunResult | None]: """Load the data from a results JSON into a dict. Parameters @@ -63,19 +63,19 @@ def _load_valid_result_json(fpath: os.PathLike | str) -> tuple[tuple | None, dic # TODO: only load this once during collection, then pass namedtuple(fname, dict) into this function # for now though, it's not the bottleneck on performance - result = load_json(fpath) + result = QuickrunResult.from_json(fpath) try: names = _get_name(result) except (ValueError, IndexError): click.secho(f"{fpath}: Missing ligand names and/or simulation type. Skipping.",err=True, fg="yellow") # fmt: skip return None, None - if result["estimate"] is None: + if result.estimate is None: click.secho(f"{fpath}: No 'estimate' found, assuming to be a failed simulation.",err=True, fg="yellow") # fmt: skip return names, None - if result["uncertainty"] is None: + if result.uncertainty is None: click.secho(f"{fpath}: No 'uncertainty' found, assuming to be a failed simulation.",err=True, fg="yellow") # fmt: skip return names, None - if all("exception" in u for u in result["unit_results"].values()): + if all("exception" in u for u in result.unit_results.values()): click.secho(f"{fpath}: Exception found in all 'unit_results', assuming to be a failed simulation.",err=True, fg="yellow") # fmt: skip return names, None return names, result @@ -110,17 +110,17 @@ def _get_legs_from_result_jsons( if name is None: # this means it couldn't find name and/or simtype continue - dgs[name]["overall"].append([result["estimate"], result["uncertainty"]]) - proto_key = [k for k in result["unit_results"].keys() if k.startswith("ProtocolUnitResult")] + dgs[name]["overall"].append([result.estimate, result.uncertainty]) + proto_key = [k for k in result.unit_results.keys() if k.startswith("ProtocolUnitResult")] for p in proto_key: - if "unit_estimate" in result["unit_results"][p]["outputs"]: - simtype = result["unit_results"][p]["outputs"]["simtype"] - dg = result["unit_results"][p]["outputs"]["unit_estimate"] - dg_error = result["unit_results"][p]["outputs"]["unit_estimate_error"] + if "unit_estimate" in result.unit_results[p]["outputs"]: + simtype = result.unit_results[p]["outputs"]["simtype"] + dg = result.unit_results[p]["outputs"]["unit_estimate"] + dg_error = result.unit_results[p]["outputs"]["unit_estimate_error"] dgs[name][simtype].append([dg, dg_error]) - if "standard_state_correction" in result["unit_results"][p]["outputs"]: - corr = result["unit_results"][p]["outputs"]["standard_state_correction"] + if "standard_state_correction" in result.unit_results[p]["outputs"]: + corr = result.unit_results[p]["outputs"]["standard_state_correction"] dgs[name]["standard_state_correction"].append([corr, 0 * unit.kilocalorie_per_mole]) else: continue diff --git a/openfecli/commands/gather_septop.py b/openfecli/commands/gather_septop.py index 9a027a757..54ba9b041 100644 --- a/openfecli/commands/gather_septop.py +++ b/openfecli/commands/gather_septop.py @@ -14,9 +14,9 @@ from openfecli.commands.gather import ( _collect_result_jsons, format_df_with_precision, - load_json, rich_print_to_stdout, ) +from openfecli.commands.quickrun import QuickrunResult def _load_valid_result_json(fpath: os.PathLike | str) -> tuple[tuple | None, dict | None]: @@ -44,19 +44,19 @@ def _load_valid_result_json(fpath: os.PathLike | str) -> tuple[tuple | None, dic # TODO: only load this once during collection, then pass namedtuple(fname, dict) into this function # for now though, it's not the bottleneck on performance - result = load_json(fpath) + result = QuickrunResult.from_json(fpath) try: names = _get_names(result) except (ValueError, IndexError): click.secho(f"{fpath}: Missing ligand names and/or simulation type. Skipping.",err=True, fg="yellow") # fmt: skip return None, None - if result["estimate"] is None: + if result.estimate is None: click.secho(f"{fpath}: No 'estimate' found, assuming to be a failed simulation.",err=True, fg="yellow") # fmt: skip return names, None - if result["uncertainty"] is None: + if result.uncertainty is None: click.secho(f"{fpath}: No 'uncertainty' found, assuming to be a failed simulation.",err=True, fg="yellow") # fmt: skip return names, None - if all("exception" in u for u in result["unit_results"].values()): + if all("exception" in u for u in result.unit_results.values()): click.secho(f"{fpath}: Exception found in all 'unit_results', assuming to be a failed simulation.",err=True, fg="yellow") # fmt: skip return names, None return names, result @@ -91,22 +91,22 @@ def _get_legs_from_result_jsons( if names is None: # this means it couldn't find names and/or simtype continue - ddgs[names]["overall"].append([result["estimate"], result["uncertainty"]]) + ddgs[names]["overall"].append([result.estimate, result.uncertainty]) proto_key = [ k - for k in result["unit_results"].keys() + for k in result.unit_results.keys() if k.startswith("ProtocolUnitResult") ] # fmt: skip for p in proto_key: - if "unit_estimate" in result["unit_results"][p]["outputs"]: - simtype = result["unit_results"][p]["outputs"]["simtype"] - dg = result["unit_results"][p]["outputs"]["unit_estimate"] - dg_error = result["unit_results"][p]["outputs"]["unit_estimate_error"] + if "unit_estimate" in result.unit_results[p]["outputs"]: + simtype = result.unit_results[p]["outputs"]["simtype"] + dg = result.unit_results[p]["outputs"]["unit_estimate"] + dg_error = result.unit_results[p]["outputs"]["unit_estimate_error"] ddgs[names][simtype].append([dg, dg_error]) - elif "standard_state_correction_A" in result["unit_results"][p]["outputs"]: - corr_A = result["unit_results"][p]["outputs"]["standard_state_correction_A"] - corr_B = result["unit_results"][p]["outputs"]["standard_state_correction_B"] + elif "standard_state_correction_A" in result.unit_results[p]["outputs"]: + corr_A = result.unit_results[p]["outputs"]["standard_state_correction_A"] + corr_B = result.unit_results[p]["outputs"]["standard_state_correction_B"] ddgs[names]["standard_state_correction_A"].append( [corr_A, 0 * unit.kilocalorie_per_mole] ) @@ -119,7 +119,7 @@ def _get_legs_from_result_jsons( return ddgs -def _get_names(result: dict) -> tuple[str, str]: +def _get_names(result: QuickrunResult) -> tuple[str, str]: """Get the ligand names from a unit's results data. Parameters @@ -133,7 +133,7 @@ def _get_names(result: dict) -> tuple[str, str]: Ligand names corresponding to the results. """ - solvent_data = list(result["protocol_result"]["data"]["solvent"].values())[0][0] + solvent_data = list(result.protocol_result["data"]["solvent"].values())[0][0] name_A = solvent_data["inputs"]["alchemical_components"]["stateA"][0]["molprops"]["ofe-name"] name_B = solvent_data["inputs"]["alchemical_components"]["stateB"][0]["molprops"]["ofe-name"] From 375899fce799f38975224705b6bcfa8e6c266055 Mon Sep 17 00:00:00 2001 From: Alyssa Travitz Date: Wed, 14 Jan 2026 10:05:42 -0800 Subject: [PATCH 4/8] move quickrun result to its own file --- openfecli/commands/gather_abfe.py | 2 +- openfecli/commands/gather_septop.py | 2 +- openfecli/commands/quickrun.py | 52 +---------------------------- openfecli/quickrun_result.py | 52 +++++++++++++++++++++++++++++ 4 files changed, 55 insertions(+), 53 deletions(-) create mode 100644 openfecli/quickrun_result.py diff --git a/openfecli/commands/gather_abfe.py b/openfecli/commands/gather_abfe.py index 26a02de9e..88c19fc4f 100644 --- a/openfecli/commands/gather_abfe.py +++ b/openfecli/commands/gather_abfe.py @@ -15,7 +15,7 @@ format_df_with_precision, rich_print_to_stdout, ) -from openfecli.commands.quickrun import QuickrunResult +from openfecli.quickrun_result import QuickrunResult def _get_name(result: QuickrunResult) -> str: diff --git a/openfecli/commands/gather_septop.py b/openfecli/commands/gather_septop.py index 54ba9b041..e168756bc 100644 --- a/openfecli/commands/gather_septop.py +++ b/openfecli/commands/gather_septop.py @@ -16,7 +16,7 @@ format_df_with_precision, rich_print_to_stdout, ) -from openfecli.commands.quickrun import QuickrunResult +from openfecli.quickrun_result import QuickrunResult def _load_valid_result_json(fpath: os.PathLike | str) -> tuple[tuple | None, dict | None]: diff --git a/openfecli/commands/quickrun.py b/openfecli/commands/quickrun.py index 9986b118e..4aa8c8bfd 100644 --- a/openfecli/commands/quickrun.py +++ b/openfecli/commands/quickrun.py @@ -1,17 +1,12 @@ # This code is part of OpenFE and is licensed under the MIT license. # For details, see https://github.com/OpenFreeEnergy/openfe -import json import pathlib -from dataclasses import asdict, dataclass -from os import PathLike -from typing import Any, Self import click -from gufe.tokenization import JSON_HANDLER -from openff.units import Quantity from openfecli import OFECommandPlugin +from openfecli.quickrun_result import QuickrunResult from openfecli.utils import configure_logger, print_duration, write @@ -20,51 +15,6 @@ def _format_exception(exception) -> str: return f"{exception[0]}: {exception[1][0]}" -@dataclass -class QuickrunResult: - """ - Class for storing protocol result data along with useful metadata. - Could ProtocolResults store this data alongside ``n_protocol_dag_results``? - """ - - estimate: Quantity - uncertainty: Quantity - protocol_result: dict[str, Any] - unit_results: dict[int, dict] - # ok: bool # TODO: add in 2.0? - - def to_json(self, filepath) -> None: - with open(filepath, mode="w") as file: - json.dump(asdict(self), file, cls=JSON_HANDLER.encoder) - - @classmethod - def from_json(cls, file: PathLike | None, content: str | None = None) -> Self: - """Load a JSON file containing a gufe object. - - Parameters - ---------- - fpath : os.PathLike | str - The path to a results JSON generated by ``openfe quickrun``. - - - Returns - ------- - QuickrunResult - A QuickrunResult instance containing the data from ``file, co``. - - """ - # similar to gufe.tokenization.from_json - if content is not None and file is not None: - raise ValueError("Cannot specify both `content` and `file`; only one input allowed") - elif content is None and file is None: - raise ValueError("Must specify either `content` and `file` for JSON input") - if file: - data = json.load(open(file, "r"), cls=JSON_HANDLER.decoder) - if content: - data = json.loads(content, cls=JSON_HANDLER.decoder) - return cls(**data) - - @click.command("quickrun", short_help="Run a given transformation, saved as a JSON file") @click.argument("transformation", type=click.File(mode="r"), required=True) @click.option( diff --git a/openfecli/quickrun_result.py b/openfecli/quickrun_result.py new file mode 100644 index 000000000..a7a0733c9 --- /dev/null +++ b/openfecli/quickrun_result.py @@ -0,0 +1,52 @@ +import json +from dataclasses import asdict, dataclass +from os import PathLike +from typing import Any, Self + +from gufe.tokenization import JSON_HANDLER +from openff.units import Quantity + + +@dataclass +class QuickrunResult: + """ + Class for storing protocol result data along with useful metadata. + Could ProtocolResults store this data alongside ``n_protocol_dag_results``? + """ + + estimate: Quantity + uncertainty: Quantity + protocol_result: dict[str, Any] + unit_results: dict[int, dict] + # ok: bool # TODO: add in 2.0? + + def to_json(self, filepath) -> None: + with open(filepath, mode="w") as file: + json.dump(asdict(self), file, cls=JSON_HANDLER.encoder) + + @classmethod + def from_json(cls, file: PathLike | None, content: str | None = None) -> Self: + """Load a JSON file containing a gufe object. + + Parameters + ---------- + fpath : os.PathLike | str + The path to a results JSON generated by ``openfe quickrun``. + + + Returns + ------- + QuickrunResult + A QuickrunResult instance containing the data from ``file, co``. + + """ + # similar to gufe.tokenization.from_json + if content is not None and file is not None: + raise ValueError("Cannot specify both `content` and `file`; only one input allowed") + elif content is None and file is None: + raise ValueError("Must specify either `content` and `file` for JSON input") + if file: + data = json.load(open(file, "r"), cls=JSON_HANDLER.decoder) + if content: + data = json.loads(content, cls=JSON_HANDLER.decoder) + return cls(**data) From e240b8633d5f8370ad17493b1bd26fcf14d4293c Mon Sep 17 00:00:00 2001 From: Alyssa Travitz Date: Wed, 14 Jan 2026 13:53:50 -0800 Subject: [PATCH 5/8] adding some ideas --- openfecli/quickrun_result.py | 26 +++++++++++++++++++++++++- 1 file changed, 25 insertions(+), 1 deletion(-) diff --git a/openfecli/quickrun_result.py b/openfecli/quickrun_result.py index a7a0733c9..b3418b949 100644 --- a/openfecli/quickrun_result.py +++ b/openfecli/quickrun_result.py @@ -1,8 +1,9 @@ import json from dataclasses import asdict, dataclass from os import PathLike -from typing import Any, Self +from typing import Any, Iterable, Self, Tuple +from gufe import ChemicalSystem from gufe.tokenization import JSON_HANDLER from openff.units import Quantity @@ -19,6 +20,8 @@ class QuickrunResult: protocol_result: dict[str, Any] unit_results: dict[int, dict] # ok: bool # TODO: add in 2.0? + # id: Tuple = () + # run_type: # protocol-specific? def to_json(self, filepath) -> None: with open(filepath, mode="w") as file: @@ -50,3 +53,24 @@ def from_json(cls, file: PathLike | None, content: str | None = None) -> Self: if content: data = json.loads(content, cls=JSON_HANDLER.decoder) return cls(**data) + + def _load_valid_result_json(self): + pass + + def run_type(self): + pass + + def stateA(self): + pass + + def stateB(self): + pass + + +@dataclass +class QuickrunResultsNetwork: + edges: Iterable[QuickrunResult] + nodes: Iterable[ChemicalSystem] # TODO: should this be a ChemicalSystem or SMC? + + def to_FEMap(self): + pass From e685ff3c38b748f778340b8659b34d67c3bfd24b Mon Sep 17 00:00:00 2001 From: Alyssa Travitz Date: Wed, 21 Jan 2026 12:09:48 -0800 Subject: [PATCH 6/8] remove boilerplate --- openfecli/quickrun_result.py | 14 +------------- 1 file changed, 1 insertion(+), 13 deletions(-) diff --git a/openfecli/quickrun_result.py b/openfecli/quickrun_result.py index b3418b949..c8645e54f 100644 --- a/openfecli/quickrun_result.py +++ b/openfecli/quickrun_result.py @@ -1,7 +1,7 @@ import json from dataclasses import asdict, dataclass from os import PathLike -from typing import Any, Iterable, Self, Tuple +from typing import Any, Self from gufe import ChemicalSystem from gufe.tokenization import JSON_HANDLER @@ -19,9 +19,6 @@ class QuickrunResult: uncertainty: Quantity protocol_result: dict[str, Any] unit_results: dict[int, dict] - # ok: bool # TODO: add in 2.0? - # id: Tuple = () - # run_type: # protocol-specific? def to_json(self, filepath) -> None: with open(filepath, mode="w") as file: @@ -65,12 +62,3 @@ def stateA(self): def stateB(self): pass - - -@dataclass -class QuickrunResultsNetwork: - edges: Iterable[QuickrunResult] - nodes: Iterable[ChemicalSystem] # TODO: should this be a ChemicalSystem or SMC? - - def to_FEMap(self): - pass From b5d7db4faee88596190995e92a5f9217fb70d619 Mon Sep 17 00:00:00 2001 From: Alyssa Travitz Date: Wed, 21 Jan 2026 12:11:29 -0800 Subject: [PATCH 7/8] make private --- openfecli/commands/gather.py | 14 ++++++++------ openfecli/commands/gather_abfe.py | 10 ++++++---- openfecli/commands/gather_septop.py | 6 +++--- openfecli/commands/quickrun.py | 5 ++--- openfecli/quickrun_result.py | 7 +++---- openfecli/tests/commands/test_gather.py | 18 +++++++++--------- 6 files changed, 31 insertions(+), 29 deletions(-) diff --git a/openfecli/commands/gather.py b/openfecli/commands/gather.py index 543bf60a9..655907a98 100644 --- a/openfecli/commands/gather.py +++ b/openfecli/commands/gather.py @@ -11,7 +11,7 @@ from openfecli import OFECommandPlugin from openfecli.clicktypes import HyphenAwareChoice -from openfecli.commands.quickrun import QuickrunResult +from openfecli.commands.quickrun import _QuickrunResult FAIL_STR = "Error" # string used to indicate a failed run in output tables. @@ -182,7 +182,7 @@ def is_results_json(fpath: os.PathLike | str) -> bool: return "estimate" in open(fpath, "r").read(20) -def _get_names(result: QuickrunResult) -> tuple[str, str]: +def _get_names(result: _QuickrunResult) -> tuple[str, str]: """Get the ligand names from a unit's results data. Parameters @@ -212,7 +212,7 @@ def _get_names(result: QuickrunResult) -> tuple[str, str]: return str(name_A), str(name_B) -def _get_type(result: QuickrunResult) -> Literal["vacuum", "solvent", "complex"]: +def _get_type(result: _QuickrunResult) -> Literal["vacuum", "solvent", "complex"]: """Determine the simulation type based on the component types.""" protocol_data = list(result.protocol_result["data"].values())[0][0] @@ -248,7 +248,7 @@ def _legacy_get_type(res_fn: os.PathLike | str) -> Literal["vacuum", "solvent", def _get_result_id( - result: QuickrunResult, result_fn: os.PathLike | str + result: _QuickrunResult, result_fn: os.PathLike | str ) -> tuple[tuple[str, str], Literal["vacuum", "solvent", "complex"]]: """Extract the name and simulation type from a results dict. @@ -274,7 +274,9 @@ def _get_result_id( return (ligA, ligB), simtype -def _load_valid_result_json(fpath: os.PathLike | str) -> tuple[tuple | None, QuickrunResult | None]: +def _load_valid_result_json( + fpath: os.PathLike | str, +) -> tuple[tuple | None, _QuickrunResult | None]: """Load the data from a results JSON into a dict. Parameters @@ -291,7 +293,7 @@ def _load_valid_result_json(fpath: os.PathLike | str) -> tuple[tuple | None, Qui """ # TODO: only load this once during collection, then pass namedtuple(fname, dict) into this function # for now though, it's not the bottleneck on performance - result = QuickrunResult.from_json(fpath) + result = _QuickrunResult.from_json(fpath) try: result_id = _get_result_id(result, fpath) diff --git a/openfecli/commands/gather_abfe.py b/openfecli/commands/gather_abfe.py index 88c19fc4f..6b836585f 100644 --- a/openfecli/commands/gather_abfe.py +++ b/openfecli/commands/gather_abfe.py @@ -15,10 +15,10 @@ format_df_with_precision, rich_print_to_stdout, ) -from openfecli.quickrun_result import QuickrunResult +from openfecli.quickrun_result import _QuickrunResult -def _get_name(result: QuickrunResult) -> str: +def _get_name(result: _QuickrunResult) -> str: """Get the ligand name from a unit's results data. Parameters @@ -38,7 +38,9 @@ def _get_name(result: QuickrunResult) -> str: return str(name) -def _load_valid_result_json(fpath: os.PathLike | str) -> tuple[tuple | None, QuickrunResult | None]: +def _load_valid_result_json( + fpath: os.PathLike | str, +) -> tuple[tuple | None, _QuickrunResult | None]: """Load the data from a results JSON into a dict. Parameters @@ -63,7 +65,7 @@ def _load_valid_result_json(fpath: os.PathLike | str) -> tuple[tuple | None, Qui # TODO: only load this once during collection, then pass namedtuple(fname, dict) into this function # for now though, it's not the bottleneck on performance - result = QuickrunResult.from_json(fpath) + result = _QuickrunResult.from_json(fpath) try: names = _get_name(result) except (ValueError, IndexError): diff --git a/openfecli/commands/gather_septop.py b/openfecli/commands/gather_septop.py index e168756bc..3fae2a171 100644 --- a/openfecli/commands/gather_septop.py +++ b/openfecli/commands/gather_septop.py @@ -16,7 +16,7 @@ format_df_with_precision, rich_print_to_stdout, ) -from openfecli.quickrun_result import QuickrunResult +from openfecli.quickrun_result import _QuickrunResult def _load_valid_result_json(fpath: os.PathLike | str) -> tuple[tuple | None, dict | None]: @@ -44,7 +44,7 @@ def _load_valid_result_json(fpath: os.PathLike | str) -> tuple[tuple | None, dic # TODO: only load this once during collection, then pass namedtuple(fname, dict) into this function # for now though, it's not the bottleneck on performance - result = QuickrunResult.from_json(fpath) + result = _QuickrunResult.from_json(fpath) try: names = _get_names(result) except (ValueError, IndexError): @@ -119,7 +119,7 @@ def _get_legs_from_result_jsons( return ddgs -def _get_names(result: QuickrunResult) -> tuple[str, str]: +def _get_names(result: _QuickrunResult) -> tuple[str, str]: """Get the ligand names from a unit's results data. Parameters diff --git a/openfecli/commands/quickrun.py b/openfecli/commands/quickrun.py index 4aa8c8bfd..f52df614e 100644 --- a/openfecli/commands/quickrun.py +++ b/openfecli/commands/quickrun.py @@ -6,7 +6,7 @@ import click from openfecli import OFECommandPlugin -from openfecli.quickrun_result import QuickrunResult +from openfecli.quickrun_result import _QuickrunResult from openfecli.utils import configure_logger, print_duration, write @@ -114,12 +114,11 @@ def quickrun(transformation, work_dir, output): else: estimate = uncertainty = None # for output file - quickrun_result = QuickrunResult( + quickrun_result = _QuickrunResult( estimate=estimate, uncertainty=uncertainty, protocol_result=prot_result.to_dict(), unit_results={unit.key: unit.to_keyed_dict() for unit in dagresult.protocol_unit_results}, - # ok=dagresult.ok() # TODO: add in 2.0 ) quickrun_result.to_json(output) diff --git a/openfecli/quickrun_result.py b/openfecli/quickrun_result.py index c8645e54f..2424991a0 100644 --- a/openfecli/quickrun_result.py +++ b/openfecli/quickrun_result.py @@ -3,13 +3,12 @@ from os import PathLike from typing import Any, Self -from gufe import ChemicalSystem from gufe.tokenization import JSON_HANDLER from openff.units import Quantity @dataclass -class QuickrunResult: +class _QuickrunResult: """ Class for storing protocol result data along with useful metadata. Could ProtocolResults store this data alongside ``n_protocol_dag_results``? @@ -36,8 +35,8 @@ def from_json(cls, file: PathLike | None, content: str | None = None) -> Self: Returns ------- - QuickrunResult - A QuickrunResult instance containing the data from ``file, co``. + _QuickrunResult + A _QuickrunResult instance containing the data from ``file, co``. """ # similar to gufe.tokenization.from_json diff --git a/openfecli/tests/commands/test_gather.py b/openfecli/tests/commands/test_gather.py index 7aa149b49..aa506f734 100644 --- a/openfecli/tests/commands/test_gather.py +++ b/openfecli/tests/commands/test_gather.py @@ -17,7 +17,7 @@ ) from openfecli.commands.gather_abfe import gather_abfe from openfecli.commands.gather_septop import gather_septop -from openfecli.commands.quickrun import QuickrunResult +from openfecli.commands.quickrun import _QuickrunResult from ..conftest import HAS_INTERNET from ..utils import assert_click_success @@ -73,13 +73,13 @@ def test_get_column(val, col): @pytest.fixture def min_valid_quickrun_result(min_result_json): - return QuickrunResult(**min_result_json) + return _QuickrunResult(**min_result_json) class TestResultLoading: def test_minimal_valid_results(self, capsys, min_valid_quickrun_result): with mock.patch( - "openfecli.commands.gather.QuickrunResult.from_json", + "openfecli.commands.gather._QuickrunResult.from_json", return_value=min_valid_quickrun_result, ): result = _load_valid_result_json(fpath="") @@ -91,7 +91,7 @@ def test_skip_missing_unit_result(self, capsys, min_valid_quickrun_result): min_valid_quickrun_result.unit_results = {} with mock.patch( - "openfecli.commands.gather.QuickrunResult.from_json", + "openfecli.commands.gather._QuickrunResult.from_json", return_value=min_valid_quickrun_result, ): result = _load_valid_result_json(fpath="") @@ -103,7 +103,7 @@ def test_skip_missing_estimate(self, capsys, min_valid_quickrun_result): min_valid_quickrun_result.estimate = None with mock.patch( - "openfecli.commands.gather.QuickrunResult.from_json", + "openfecli.commands.gather._QuickrunResult.from_json", return_value=min_valid_quickrun_result, ): result = _load_valid_result_json(fpath="") @@ -115,7 +115,7 @@ def test_skip_missing_uncertainty(self, capsys, min_valid_quickrun_result): min_valid_quickrun_result.uncertainty = None with mock.patch( - "openfecli.commands.gather.QuickrunResult.from_json", + "openfecli.commands.gather._QuickrunResult.from_json", return_value=min_valid_quickrun_result, ): result = _load_valid_result_json(fpath="") @@ -126,7 +126,7 @@ def test_skip_missing_uncertainty(self, capsys, min_valid_quickrun_result): def test_skip_all_failed_runs(self, capsys, min_valid_quickrun_result): del min_valid_quickrun_result.unit_results["ProtocolUnitResult-e85"] with mock.patch( - "openfecli.commands.gather.QuickrunResult.from_json", + "openfecli.commands.gather._QuickrunResult.from_json", return_value=min_valid_quickrun_result, ): result = _load_valid_result_json(fpath="") @@ -137,7 +137,7 @@ def test_skip_all_failed_runs(self, capsys, min_valid_quickrun_result): def test_missing_pr_data(self, capsys, min_valid_quickrun_result): min_valid_quickrun_result.protocol_result["data"] = {} with mock.patch( - "openfecli.commands.gather.QuickrunResult.from_json", + "openfecli.commands.gather._QuickrunResult.from_json", return_value=min_valid_quickrun_result, ): result = _load_valid_result_json(fpath="") @@ -150,7 +150,7 @@ def test_get_legs_from_result_jsons(self, capsys, min_valid_quickrun_result): min_valid_quickrun_result.protocol_result["data"] = {} with mock.patch( - "openfecli.commands.gather.QuickrunResult.from_json", + "openfecli.commands.gather._QuickrunResult.from_json", return_value=min_valid_quickrun_result, ): result = _get_legs_from_result_jsons(result_fns=[""], report="dg") From ebbbbf8caecb406063710e9c28f0b0aeef93df5f Mon Sep 17 00:00:00 2001 From: Alyssa Travitz Date: Wed, 21 Jan 2026 14:12:46 -0800 Subject: [PATCH 8/8] remove unused code --- openfecli/quickrun_result.py | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/openfecli/quickrun_result.py b/openfecli/quickrun_result.py index 2424991a0..e3c79118b 100644 --- a/openfecli/quickrun_result.py +++ b/openfecli/quickrun_result.py @@ -49,15 +49,3 @@ def from_json(cls, file: PathLike | None, content: str | None = None) -> Self: if content: data = json.loads(content, cls=JSON_HANDLER.decoder) return cls(**data) - - def _load_valid_result_json(self): - pass - - def run_type(self): - pass - - def stateA(self): - pass - - def stateB(self): - pass