Issue 695 export variables

.gitignore

@@ -5,11 +5,13 @@
 /local/
 *.DS_Store
 *.mat
+*.csv
 
-# don't ignore important .txt files
+# don't ignore important .txt and .csv files
 !requirements*
 !LICENSE.txt
 !CMakeLists.txt
+!input/**/*.csv
 
 # running files
 *.pyc

examples/scripts/experimental_protocols/cccv.py

@@ -16,7 +16,7 @@
     * 3,
     period="2 minutes",
 )
-model = pybamm.lithium_ion.DFN()
+model = pybamm.lithium_ion.DFN()  # use {"thermal": "x-lumped"} for thermal effects
 sim = pybamm.Simulation(model, experiment=experiment, solver=pybamm.CasadiSolver())
 sim.solve()
 
@@ -35,5 +35,30 @@
     ax.set_xlim([0, 13])
 ax.legend()
 
+# Save time, voltage, current, discharge capacity and temperature to csv and matlab
+# formats
+sim.solution.save_data(
+    "output.mat",
+    [
+        "Time [h]",
+        "Current [A]",
+        "Terminal voltage [V]",
+        "Discharge capacity [A.h]",
+        "X-averaged cell temperature [K]",
+    ],
+    to_format="matlab",
+)
+sim.solution.save_data(
+    "output.csv",
+    [
+        "Time [h]",
+        "Current [A]",
+        "Terminal voltage [V]",
+        "Discharge capacity [A.h]",
+        "X-averaged cell temperature [K]",
+    ],
+    to_format="csv",
+)
+
 # Show all plots
 sim.plot()

pybamm/solvers/solution.py

@@ -6,7 +6,9 @@
 import numpy as np
 import pickle
 import pybamm
+import pandas as pd
 from collections import defaultdict
+from scipy.io import savemat
 
 
 class _BaseSolution(object):
@@ -190,17 +192,56 @@ def save(self, filename):
         with open(filename, "wb") as f:
             pickle.dump(self, f, pickle.HIGHEST_PROTOCOL)
 
-    def save_data(self, filename):
-        """Save solution data only (raw arrays) using pickle"""
-        if len(self.data) == 0:
+    def save_data(self, filename, variables=None, to_format="pickle"):
+        """
+        Save solution data only (raw arrays)
+
+        Parameters
+        ----------
+        filename : str
+            The name of the file to save data to
+        variables : list, optional
+            List of variables to save. If None, saves all of the variables that have
+            been created so far
+        to_format : str, optional
+            The format to save to. Options are:
+
+            - 'pickle' (default): creates a pickle file with the data dictionary
+            - 'matlab': creates a .mat file, for loading in matlab
+            - 'csv': creates a csv file (1D variables only)
+
+        """
+        if variables is None:
+            # variables not explicitly provided -> save all variables that have been
+            # computed
+            data = self.data
+        else:
+            # otherwise, save only the variables specified
+            data = {}
+            for name in variables:
+                data[name] = self[name].data
+        if len(data) == 0:
             raise ValueError(
-                """Solution does not have any data. Add variables by calling
-                'solution.update', e.g.
-                'solution.update(["Terminal voltage [V]", "Current [A]"])'
-                and then save"""
+                """
+                Solution does not have any data. Please provide a list of variables
+                to save.
+                """
             )
-        with open(filename, "wb") as f:
-            pickle.dump(self.data, f, pickle.HIGHEST_PROTOCOL)
+        if to_format == "pickle":
+            with open(filename, "wb") as f:
+                pickle.dump(data, f, pickle.HIGHEST_PROTOCOL)
+        elif to_format == "matlab":
+            savemat(filename, data)
+        elif to_format == "csv":
+            for name, var in data.items():
+                if var.ndim == 2:
+                    raise ValueError(
+                        "only 1D variables can be saved to csv, but '{}' is 2D".format(
+                            name
+                        )
+                    )
+            df = pd.DataFrame(data)
+            df.to_csv(filename, index=False)
 
 
 class Solution(_BaseSolution):

tests/unit/test_solvers/test_solution.py

@@ -4,6 +4,9 @@
 import pybamm
 import unittest
 import numpy as np
+import pandas as pd
+from scipy.io import loadmat
+from tests import get_discretisation_for_testing
 
 
 class TestSolution(unittest.TestCase):
@@ -93,29 +96,51 @@ def test_getitem(self):
 
     def test_save(self):
         model = pybamm.BaseModel()
+        # create both 1D and 2D variables
         c = pybamm.Variable("c")
-        model.rhs = {c: -c}
-        model.initial_conditions = {c: 1}
-        model.variables["c"] = c
+        d = pybamm.Variable("d", domain="negative electrode")
+        model.rhs = {c: -c, d: 1}
+        model.initial_conditions = {c: 1, d: 2}
+        model.variables = {"c": c, "d": d, "2c": 2 * c}
 
-        disc = pybamm.Discretisation()
+        disc = get_discretisation_for_testing()
         disc.process_model(model)
         solution = pybamm.ScipySolver().solve(model, np.linspace(0, 1))
 
         # test save data
         with self.assertRaises(ValueError):
             solution.save_data("test.pickle")
         # set variables first then save
-        solution.update(["c"])
+        solution.update(["c", "d"])
         solution.save_data("test.pickle")
         data_load = pybamm.load("test.pickle")
         np.testing.assert_array_equal(solution.data["c"], data_load["c"])
-
-        # test save
+        np.testing.assert_array_equal(solution.data["d"], data_load["d"])
+
+        # to matlab
+        solution.save_data("test.mat", to_format="matlab")
+        data_load = loadmat("test.mat")
+        np.testing.assert_array_equal(solution.data["c"], data_load["c"].flatten())
+        np.testing.assert_array_equal(solution.data["d"], data_load["d"])
+
+        # to csv
+        with self.assertRaisesRegex(
+            ValueError, "only 1D variables can be saved to csv"
+        ):
+            solution.save_data("test.csv", to_format="csv")
+        # only save "c" and "2c"
+        solution.save_data("test.csv", ["c", "2c"], to_format="csv")
+        # 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"])
+
+        # test save whole solution
         solution.save("test.pickle")
         solution_load = pybamm.load("test.pickle")
         self.assertEqual(solution.model.name, solution_load.model.name)
         np.testing.assert_array_equal(solution["c"].entries, solution_load["c"].entries)
+        np.testing.assert_array_equal(solution["d"].entries, solution_load["d"].entries)
 
     def test_solution_evals_with_inputs(self):
         model = pybamm.lithium_ion.SPM()