Adds code to save and load results objects (#161)

DrPaulSharp · web-flow · commit b9a7b1071381 · 2025-05-14T17:49:24.000+01:00
* Adds code to save and load results objects

* Merges load methods

* Addresses review comments
diff --git a/RATapi/__init__.py b/RATapi/__init__.py
@@ -6,11 +6,24 @@
 from RATapi import events, models
 from RATapi.classlist import ClassList
 from RATapi.controls import Controls
+from RATapi.outputs import BayesResults, Results
 from RATapi.project import Project
 from RATapi.run import run
 from RATapi.utils import convert, plotting
 
 with suppress(ImportError):  # orsopy is an optional dependency
     from RATapi.utils import orso as orso
 
-__all__ = ["examples", "models", "events", "ClassList", "Controls", "Project", "run", "plotting", "convert"]
+__all__ = [
+    "examples",
+    "models",
+    "events",
+    "ClassList",
+    "Controls",
+    "BayesResults",
+    "Results",
+    "Project",
+    "run",
+    "plotting",
+    "convert",
+]
diff --git a/RATapi/outputs.py b/RATapi/outputs.py
@@ -1,13 +1,78 @@
 """Converts results from the compiled RAT code to python dataclasses."""
 
+import json
 from dataclasses import dataclass
+from pathlib import Path
 from typing import Any, Optional, Union
 
 import numpy as np
 
 import RATapi.rat_core
 from RATapi.utils.enums import Procedures
 
+bayes_results_subclasses = [
+    "predictionIntervals",
+    "confidenceIntervals",
+    "dreamParams",
+    "dreamOutput",
+    "nestedSamplerOutput",
+]
+
+bayes_results_fields = {
+    "param_fields": {
+        "predictionIntervals": [],
+        "confidenceIntervals": [],
+        "dreamParams": [
+            "nParams",
+            "nChains",
+            "nGenerations",
+            "parallel",
+            "CPU",
+            "jumpProbability",
+            "pUnitGamma",
+            "nCR",
+            "delta",
+            "steps",
+            "zeta",
+            "outlier",
+            "adaptPCR",
+            "thinning",
+            "epsilon",
+            "ABC",
+            "IO",
+            "storeOutput",
+        ],
+        "dreamOutput": ["runtime", "iteration"],
+        "nestedSamplerOutput": ["logZ", "logZErr"],
+    },
+    "list_fields": {
+        "predictionIntervals": ["reflectivity"],
+        "confidenceIntervals": [],
+        "dreamParams": [],
+        "dreamOutput": [],
+        "nestedSamplerOutput": [],
+    },
+    "double_list_fields": {
+        "predictionIntervals": ["sld"],
+        "confidenceIntervals": [],
+        "dreamParams": [],
+        "dreamOutput": [],
+        "nestedSamplerOutput": [],
+    },
+    "array_fields": {
+        "predictionIntervals": ["sampleChi"],
+        "confidenceIntervals": ["percentile65", "percentile95", "mean"],
+        "dreamParams": ["R"],
+        "dreamOutput": ["allChains", "outlierChains", "AR", "R_stat", "CR"],
+        "nestedSamplerOutput": ["nestSamples", "postSamples"],
+    },
+}
+
+results_fields = {
+    "list_fields": ["reflectivity", "simulation", "shiftedData", "backgrounds", "resolutions"],
+    "double_list_fields": ["sldProfiles", "layers", "resampledLayers"],
+}
+
 
 def get_field_string(field: str, value: Any, array_limit: int):
     """Return a string representation of class fields where large arrays are represented by their shape.
@@ -179,6 +244,74 @@ def __str__(self):
             output += get_field_string(key, value, 100)
         return output
 
+    def save(self, filepath: Union[str, Path] = "./results.json"):
+        """Save the Results object to a JSON file.
+
+        Parameters
+        ----------
+        filepath : str or Path
+            The path to where the results file will be written.
+        """
+        filepath = Path(filepath).with_suffix(".json")
+        json_dict = write_core_results_fields(self)
+
+        filepath.write_text(json.dumps(json_dict))
+
+    @classmethod
+    def load(cls, path: Union[str, Path]) -> Union["Results", "BayesResults"]:
+        """Load a Results object from file.
+
+        Parameters
+        ----------
+        path : str or Path
+            The path to the results json file.
+        """
+        path = Path(path)
+        input_data = path.read_text()
+        results_dict = json.loads(input_data)
+
+        results_dict = read_core_results_fields(results_dict)
+
+        if all(key in results_dict for key in bayes_results_subclasses):
+            results_dict = read_bayes_results_fields(results_dict)
+
+            return BayesResults(
+                reflectivity=results_dict["reflectivity"],
+                simulation=results_dict["simulation"],
+                shiftedData=results_dict["shiftedData"],
+                backgrounds=results_dict["backgrounds"],
+                resolutions=results_dict["resolutions"],
+                sldProfiles=results_dict["sldProfiles"],
+                layers=results_dict["layers"],
+                resampledLayers=results_dict["resampledLayers"],
+                calculationResults=CalculationResults(**results_dict["calculationResults"]),
+                contrastParams=ContrastParams(**results_dict["contrastParams"]),
+                fitParams=np.array(results_dict["fitParams"]),
+                fitNames=results_dict["fitNames"],
+                predictionIntervals=PredictionIntervals(**results_dict["predictionIntervals"]),
+                confidenceIntervals=ConfidenceIntervals(**results_dict["confidenceIntervals"]),
+                dreamParams=DreamParams(**results_dict["dreamParams"]),
+                dreamOutput=DreamOutput(**results_dict["dreamOutput"]),
+                nestedSamplerOutput=NestedSamplerOutput(**results_dict["nestedSamplerOutput"]),
+                chain=np.array(results_dict["chain"]),
+            )
+
+        else:
+            return Results(
+                reflectivity=results_dict["reflectivity"],
+                simulation=results_dict["simulation"],
+                shiftedData=results_dict["shiftedData"],
+                backgrounds=results_dict["backgrounds"],
+                resolutions=results_dict["resolutions"],
+                sldProfiles=results_dict["sldProfiles"],
+                layers=results_dict["layers"],
+                resampledLayers=results_dict["resampledLayers"],
+                calculationResults=CalculationResults(**results_dict["calculationResults"]),
+                contrastParams=ContrastParams(**results_dict["contrastParams"]),
+                fitParams=np.array(results_dict["fitParams"]),
+                fitNames=results_dict["fitNames"],
+            )
+
 
 @dataclass
 class PredictionIntervals(RATResult):
@@ -405,6 +538,148 @@ class BayesResults(Results):
     nestedSamplerOutput: NestedSamplerOutput
     chain: np.ndarray
 
+    def save(self, filepath: Union[str, Path] = "./results.json"):
+        """Save the BayesResults object to a JSON file.
+
+        Parameters
+        ----------
+        filepath : str or Path
+            The path to where the results file will be written.
+        """
+        filepath = Path(filepath).with_suffix(".json")
+        json_dict = write_core_results_fields(self)
+
+        # Take each of the subclasses in a BayesResults instance and switch the numpy arrays to lists
+        for subclass_name in bayes_results_subclasses:
+            subclass = getattr(self, subclass_name)
+            subclass_dict = {}
+
+            for field in bayes_results_fields["param_fields"][subclass_name]:
+                subclass_dict[field] = getattr(subclass, field)
+
+            for field in bayes_results_fields["list_fields"][subclass_name]:
+                subclass_dict[field] = [result_array.tolist() for result_array in getattr(subclass, field)]
+
+            for field in bayes_results_fields["double_list_fields"][subclass_name]:
+                subclass_dict[field] = [
+                    [result_array.tolist() for result_array in inner_list] for inner_list in getattr(subclass, field)
+                ]
+
+            for field in bayes_results_fields["array_fields"][subclass_name]:
+                subclass_dict[field] = getattr(subclass, field).tolist()
+
+            json_dict[subclass_name] = subclass_dict
+
+        json_dict["chain"] = self.chain.tolist()
+        filepath.write_text(json.dumps(json_dict))
+
+
+def write_core_results_fields(results: Union[Results, BayesResults], json_dict: Optional[dict] = None) -> dict:
+    """Modify the values of the fields that appear in both Results and BayesResults when saving to a json file.
+
+    Parameters
+    ----------
+    results: Union[Results, BayesResults]
+        The results or BayesResults object we are writing to json.
+    json_dict: Optional[dict]
+        The dictionary containing the json output.
+
+    Returns
+    -------
+    json_dict: dict
+        The output json dict updated with the fields that appear in both Results and BayesResults.
+    """
+    if json_dict is None:
+        json_dict = {}
+
+    for field in results_fields["list_fields"]:
+        json_dict[field] = [result_array.tolist() for result_array in getattr(results, field)]
+
+    for field in results_fields["double_list_fields"]:
+        json_dict[field] = [
+            [result_array.tolist() for result_array in inner_list] for inner_list in getattr(results, field)
+        ]
+
+    json_dict["calculationResults"] = {}
+    json_dict["calculationResults"]["chiValues"] = results.calculationResults.chiValues.tolist()
+    json_dict["calculationResults"]["sumChi"] = results.calculationResults.sumChi
+
+    json_dict["contrastParams"] = {}
+    for field in results.contrastParams.__dict__:
+        json_dict["contrastParams"][field] = getattr(results.contrastParams, field).tolist()
+
+    json_dict["fitParams"] = results.fitParams.tolist()
+    json_dict["fitNames"] = results.fitNames
+
+    return json_dict
+
+
+def read_core_results_fields(results_dict: dict) -> dict:
+    """Modify the values of the fields that appear in both Results and BayesResults when loading a json file.
+
+    Parameters
+    ----------
+    results_dict: Optional[dict]
+        The dictionary containing the json input.
+
+    Returns
+    -------
+    results_dict: dict
+        The input json dict with the fields that appear in both Results and BayesResults converted to numpy arrays
+        where necessary.
+    """
+    for field in results_fields["list_fields"]:
+        results_dict[field] = [np.array(result_array) for result_array in results_dict[field]]
+
+    for field in results_fields["double_list_fields"]:
+        results_dict[field] = [
+            [np.array(result_array) for result_array in inner_list] for inner_list in results_dict[field]
+        ]
+
+    results_dict["calculationResults"]["chiValues"] = np.array(results_dict["calculationResults"]["chiValues"])
+
+    for field in results_dict["contrastParams"]:
+        results_dict["contrastParams"][field] = np.array(results_dict["contrastParams"][field])
+
+    return results_dict
+
+
+def read_bayes_results_fields(results_dict: dict) -> dict:
+    """Modify the values of the fields that appear only in BayesResults when loading a json file.
+
+    Parameters
+    ----------
+    results_dict: Optional[dict]
+        The dictionary containing the json input.
+
+    Returns
+    -------
+    results_dict: dict
+        The input json dict with the fields that appear in both Results and BayesResults converted to numpy arrays
+        where necessary.
+    """
+    for subclass_name in bayes_results_subclasses:
+        subclass_dict = {}
+
+        for field in bayes_results_fields["param_fields"][subclass_name]:
+            subclass_dict[field] = results_dict[subclass_name][field]
+
+        for field in bayes_results_fields["list_fields"][subclass_name]:
+            subclass_dict[field] = [np.array(result_array) for result_array in results_dict[subclass_name][field]]
+
+        for field in bayes_results_fields["double_list_fields"][subclass_name]:
+            subclass_dict[field] = [
+                [np.array(result_array) for result_array in inner_list]
+                for inner_list in results_dict[subclass_name][field]
+            ]
+
+        for field in bayes_results_fields["array_fields"][subclass_name]:
+            subclass_dict[field] = np.array(results_dict[subclass_name][field])
+
+        results_dict[subclass_name] = subclass_dict
+
+    return results_dict
+
 
 def make_results(
     procedure: Procedures,
diff --git a/RATapi/project.py b/RATapi/project.py
@@ -906,7 +906,6 @@ def save(self, filepath: Union[str, Path] = "./project.json"):
         ----------
         filepath : str or Path
             The path to where the project file will be written.
-
         """
         filepath = Path(filepath).with_suffix(".json")
 
diff --git a/tests/test_outputs.py b/tests/test_outputs.py
@@ -3,6 +3,9 @@
 We use the example for both a reflectivity calculation, and Bayesian analysis using the Dream algorithm.
 """
 
+import tempfile
+from pathlib import Path
+
 import numpy as np
 import pytest
 
@@ -203,3 +206,18 @@ def test_results_str(test_output_results, test_str, request) -> None:
     test_str = request.getfixturevalue(test_str)
 
     assert test_output_results.__str__() == test_str
+
+
+@pytest.mark.parametrize("result_class", [RATapi.Results, RATapi.BayesResults])
+@pytest.mark.parametrize("test_results", ["reflectivity_calculation_results", "dream_results"])
+def test_save_load(result_class, test_results, request):
+    """Test that saving and loading an output object returns the same object."""
+    test_results = request.getfixturevalue(test_results)
+
+    with tempfile.TemporaryDirectory() as tmp:
+        # ignore relative path warnings
+        path = Path(tmp, "results.json")
+        test_results.save(path)
+        loaded_results = result_class.load(path)
+
+    check_results_equal(test_results, loaded_results)
diff --git a/tests/utils.py b/tests/utils.py
