Repository: pybamm-team/PyBaMM

Issue:
save a pybamm solution object as json
**Summary**

 Extend the `pybamm.Solution.save` function to be able to export a solution in json format

 **Motivation**
 
Like csv, json is a standard recognisable format for serialising data. Unlike csv, it is more suitable for serializing 1D or 2D variables

**Additional context**
 
The save function can already output in pickle, matlab or csv format


Discussion:

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 56b62db25..7f66ed40b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -3,6 +3,7 @@
 ## Features
 
 -   Added an option to force install compatible versions of jax and jaxlib if already installed using CLI ([#1881](https://github.com/pybamm-team/PyBaMM/pull/1881))
+-   Allow pybamm.Solution.save_data() to return a string if filename is None, and added json to_format option ([#1909](https://github.com/pybamm-team/PyBaMM/pull/1909)
 
 ## Bug fixes
 
diff --git a/pybamm/solvers/solution.py b/pybamm/solvers/solution.py
index 46056ff2f..9ec9a552b 100644
--- a/pybamm/solvers/solution.py
+++ b/pybamm/solvers/solution.py
@@ -2,6 +2,7 @@
 # Solution class
 #
 import casadi
+import json
 import numbers
 import numpy as np
 import pickle
@@ -10,6 +11,19 @@ import pandas as pd
 from scipy.io import savemat
 
 
+class NumpyEncoder(json.JSONEncoder):
+    """
+    Numpy serialiser helper class that converts numpy arrays to a list
+    https://stackoverflow.com/questions/26646362/numpy-array-is-not-json-serializable
+    """
+
+    def default(self, obj):
+        if isinstance(obj, np.ndarray):
+            return obj.tolist()
+        # won't be called since we only need to convert numpy arrays
+        return json.JSONEncoder.default(self, obj)  # pragma: no cover
+
+
 class Solution(object):
     """
     Class containing the solution of, and various attributes associated with, a PyBaMM
@@ -536,14 +550,17 @@ class Solution(object):
         with open(filename, "wb") as f:
             pickle.dump(self, f, pickle.HIGHEST_PROTOCOL)
 
-    def save_data(self, filename, variables=None, to_format="pickle", short_names=None):
+    def save_data(
+            self, filename=None, variables=None,
+            to_format="pickle", short_names=None
+    ):
         """
         Save solution data only (raw arrays)
 
         Parameters
         ----------
-        filename : str
-            The name of the file to save data to
+        filename : str, optional
+            The name of the file to save data to. If None, then a str is returned
         variables : list, optional
             List of variables to save. If None, saves all of the variables that have
             been created so far
@@ -553,12 +570,19 @@ class Solution(object):
             - 'pickle' (default): creates a pickle file with the data dictionary
             - 'matlab': creates a .mat file, for loading in matlab
             - 'csv': creates a csv file (0D variables only)
+            - 'json': creates a json file
         short_names : dict, optional
             Dictionary of shortened names to use when saving. This may be necessary when
             saving to MATLAB, since no spaces or special characters are allowed in
             MATLAB variable names. Note that not all the variables need to be given
             a short name.
 
+        Returns
+        -------
+        data : str, optional
+            str if 'csv' or 'json' is chosen and filename is None, otherwise None
+
+
         """
         if variables is None:
             # variables not explicitly provided -> save all variables that have been
@@ -588,9 +612,17 @@ class Solution(object):
                 data_short_names[name] = var
 
         if to_format == "pickle":
+            if filename is None:
+                raise ValueError(
+                    "pickle format must be written to a file"
+                )
             with open(filename, "wb") as f:
                 pickle.dump(data_short_names, f, pickle.HIGHEST_PROTOCOL)
         elif to_format == "matlab":
+            if filename is None:
+                raise ValueError(
+                    "matlab format must be written to a file"
+                )
             # Check all the variable names only contain a-z, A-Z or _ or numbers
             for name in data_short_names.keys():
                 # Check the string only contains the following ASCII:
@@ -625,7 +657,13 @@ class Solution(object):
                         )
                     )
             df = pd.DataFrame(data_short_names)
-            df.to_csv(filename, index=False)
+            return df.to_csv(filename, index=False)
+        elif to_format == "json":
+            if filename is None:
+                return json.dumps(data_short_names, cls=NumpyEncoder)
+            else:
+                with open(filename, "w") as outfile:
+                    json.dump(data_short_names, outfile, cls=NumpyEncoder)
         else:
             raise ValueError("format '{}' not recognised".format(to_format))
 


Test patch:
diff --git a/tests/unit/test_solvers/test_solution.py b/tests/unit/test_solvers/test_solution.py
index 7298de45e..244c70bd3 100644
--- a/tests/unit/test_solvers/test_solution.py
+++ b/tests/unit/test_solvers/test_solution.py
@@ -1,6 +1,7 @@
 #
 # Tests for the Solution class
 #
+import json
 import pybamm
 import unittest
 import numpy as np
@@ -237,9 +238,13 @@ class TestSolution(unittest.TestCase):
         # test save data
         with self.assertRaises(ValueError):
             solution.save_data("test.pickle")
+
         # set variables first then save
         solution.update(["c", "d"])
+        with self.assertRaisesRegex(ValueError, "pickle"):
+            solution.save_data(to_format="pickle")
         solution.save_data("test.pickle")
+
         data_load = pybamm.load("test.pickle")
         np.testing.assert_array_equal(solution.data["c"], data_load["c"])
         np.testing.assert_array_equal(solution.data["d"], data_load["d"])
@@ -250,6 +255,9 @@ class TestSolution(unittest.TestCase):
         np.testing.assert_array_equal(solution.data["c"], data_load["c"].flatten())
         np.testing.assert_array_equal(solution.data["d"], data_load["d"])
 
+        with self.assertRaisesRegex(ValueError, "matlab"):
+            solution.save_data(to_format="matlab")
+
         # to matlab with bad variables name fails
         solution.update(["c + d"])
         with self.assertRaisesRegex(ValueError, "Invalid character"):
@@ -268,11 +276,36 @@ class TestSolution(unittest.TestCase):
             solution.save_data("test.csv", to_format="csv")
         # only save "c" and "2c"
         solution.save_data("test.csv", ["c", "2c"], to_format="csv")
+        csv_str = solution.save_data(variables=["c", "2c"], to_format="csv")
+
+        # check string is the same as the file
+        with open('test.csv') as f:
+            # need to strip \r chars for windows
+            self.assertEqual(
+                csv_str.replace('\r', ''), f.read()
+            )
+
         # read csv
         df = pd.read_csv("test.csv")
         np.testing.assert_array_almost_equal(df["c"], solution.data["c"])
         np.testing.assert_array_almost_equal(df["2c"], solution.data["2c"])
 
+        # to json
+        solution.save_data("test.json", to_format="json")
+        json_str = solution.save_data(to_format="json")
+
+        # check string is the same as the file
+        with open('test.json') as f:
+            # need to strip \r chars for windows
+            self.assertEqual(
+                json_str.replace('\r', ''), f.read()
+            )
+
+        # check if string has the right values
+        json_data = json.loads(json_str)
+        np.testing.assert_array_almost_equal(json_data["c"], solution.data["c"])
+        np.testing.assert_array_almost_equal(json_data["d"], solution.data["d"])
+
         # raise error if format is unknown
         with self.assertRaisesRegex(ValueError, "format 'wrong_format' not recognised"):
             solution.save_data("test.csv", to_format="wrong_format")