Adds code to save and load results objects

Author: DrPaulSharp

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",
+]

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,49 @@ 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]) -> "Results":
+        """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)
+
+        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 +513,143 @@ 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))
+
+    @classmethod
+    def load(cls, path: Union[str, Path]) -> "BayesResults":
+        """Load a BayesResults 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)
+
+        # Take each of the subclasses in a BayesResults instance and convert 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 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"]),
+        )
+
+
+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."""
+    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."""
+    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 make_results(
     procedure: Procedures,

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")
 

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,23 @@ 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", "test_results"],
+    [
+        (RATapi.Results, "reflectivity_calculation_results"),
+        (RATapi.BayesResults, "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)