Merge pull request #5177 from medha-14/parameter_serialisation

[GSoC 2025] Serialising Parameter Sets

Author: valentinsulzer

docs/conf.py

@@ -92,6 +92,7 @@
     "https://doi.org/10.1137/20M1336898",  # DOI link to ignore
     "https://en.wikipedia.org/wiki/",  # Wikipedia link to ignore
     "https://books.google.co.uk/books",
+    "https://docs.scipy.org/doc/scipy",  # SciPy docs timeout intermittently
 ]
 
 

src/pybamm/__init__.py

@@ -58,7 +58,7 @@
 from .expression_tree.operations.jacobian import Jacobian
 from .expression_tree.operations.convert_to_casadi import CasadiConverter
 from .expression_tree.operations.unpack_symbols import SymbolUnpacker
-from .expression_tree.operations.serialise import Serialise
+from .expression_tree.operations.serialise import Serialise,ExpressionFunctionParameter
 
 # Model classes
 from .models.base_model import BaseModel

src/pybamm/expression_tree/operations/serialise.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import importlib
+import inspect
 import json
 import numbers
 import re
@@ -16,6 +17,31 @@
 SUPPORTED_SCHEMA_VERSION = "1.0"
 
 
+class ExpressionFunctionParameter(pybamm.UnaryOperator):
+    def __init__(self, name, child, func_name, func_args):
+        super().__init__(name, child)
+        self.func_name = func_name
+        self.func_args = func_args
+
+    def _unary_evaluate(self, child):
+        """Evaluate the symbolic expression (the child)"""
+        return child
+
+    def to_source(self):
+        """
+        Creates python source code for the function.
+        """
+        src = f"def {self.func_name}({', '.join(self.func_args)}):\n"
+
+        expression = self.child.create_copy()
+        for child in expression.pre_order():
+            if isinstance(child, pybamm.Parameter) and child.name not in self.func_args:
+                child.name = f'Parameter("{child.name}")'
+
+        src += f"    return {expression.to_equation()}\n"
+        return src
+
+
 class Serialise:
     """
     Converts a discretised model to and from a JSON file.
@@ -1285,6 +1311,72 @@ def load_custom_model(filename: str | dict) -> pybamm.BaseModel:
 
         return model
 
+    @staticmethod
+    def save_parameters(parameters: dict, filename=None):
+        """
+        Serializes a dictionary of parameters to a JSON file.
+        The values can be numbers, PyBaMM symbols, or callables.
+
+        Parameters
+        ----------
+        parameters : dict
+            A dictionary of parameter names and values.
+            Values can be numeric, PyBaMM symbols, or callables.
+
+        filename : str, optional
+            If given, saves the serialized parameters to this file.
+        """
+        parameter_values_dict = {}
+
+        for k, v in parameters.items():
+            if callable(v):
+                parameter_values_dict[k] = Serialise.convert_symbol_to_json(
+                    Serialise.convert_function_to_symbolic_expression(v, k)
+                )
+            else:
+                parameter_values_dict[k] = Serialise.convert_symbol_to_json(v)
+
+        if filename is not None:
+            with open(filename, "w") as f:
+                json.dump(parameter_values_dict, f, indent=4)
+
+    @staticmethod
+    def load_parameters(filename):
+        """
+        Load a JSON file of parameters (either from Serialise.save_parameters
+        or from a standard pybamm.ParameterValues.save), and return a
+        pybamm.ParameterValues object.
+
+        - If a value is a dict with a "type" key, deserialize it as a PyBaMM symbol.
+        - Otherwise (float, int, bool, str, list, dict-without-type), leave it as-is.
+        """
+        with open(filename) as f:
+            raw_dict = json.load(f)
+
+        deserialized = {}
+        for key, val in raw_dict.items():
+            if isinstance(val, dict) and "type" in val:
+                deserialized[key] = Serialise.convert_symbol_from_json(val)
+
+            elif isinstance(val, list):
+                deserialized[key] = val
+
+            elif isinstance(val, (numbers.Number | bool)):
+                deserialized[key] = val
+
+            elif isinstance(val, str):
+                deserialized[key] = val
+
+            elif isinstance(val, dict):
+                deserialized[key] = val
+
+            else:
+                raise ValueError(
+                    f"Unsupported parameter format for key '{key}': {val!r}"
+                )
+
+        return pybamm.ParameterValues(deserialized)
+
     # Helper functions
 
     def _get_pybamm_class(self, snippet: dict):
@@ -1448,6 +1540,40 @@ def _convert_options(self, d):
         else:
             return d
 
+    def convert_function_to_symbolic_expression(func, name=None):
+        """
+        Converts a Python function to a PyBaMM ExpressionFunctionParameter object.
+
+
+        Parameters
+        ----------
+        func : callable
+            The Python function to convert. Its body should operate on symbolic inputs
+            (e.g., x+1, x*y) so it can be represented as a PyBaMM expression.
+
+
+        name : str, optional
+            The name of the function to use in the symbolic expression. If not provided,
+            the function's __name__ is used.
+
+
+        Returns
+        -------
+        ExpressionFunctionParameter
+            A symbolic wrapper for the function that preserves its name, arguments,
+            and expression body.
+        """
+        func_name = name or func.__name__
+
+        sig = inspect.signature(func)
+        arg_names = list(sig.parameters.keys())
+
+        sym_inputs = [pybamm.Parameter(arg) for arg in arg_names]
+
+        sym_output = func(*sym_inputs)
+
+        return ExpressionFunctionParameter(func_name, sym_output, func_name, arg_names)
+
     @staticmethod
     def convert_symbol_to_json(
         symbol: pybamm.Symbol | numbers.Number | list,
@@ -1481,8 +1607,15 @@ def convert_symbol_to_json(
         >>> Serialise.convert_symbol_to_json(v)
         {'type': 'Variable', 'name': 'c', 'domains': {'primary': [], 'secondary': [], 'tertiary': [], 'quaternary': []}, 'bounds': [{'type': 'Scalar', 'value': np.float64(-inf)}, {'type': 'Scalar', 'value': np.float64(inf)}]}
         """
-
-        if isinstance(symbol, numbers.Number | list):
+        if isinstance(symbol, ExpressionFunctionParameter):
+            return {
+                "type": "ExpressionFunctionParameter",
+                "name": symbol.name,
+                "children": [Serialise.convert_symbol_to_json(symbol.child)],
+                "func_name": symbol.func_name,
+                "func_args": symbol.func_args,
+            }
+        elif isinstance(symbol, numbers.Number | list):
             return symbol
         elif isinstance(symbol, pybamm.Time):
             return {"type": "Time"}
@@ -1692,6 +1825,13 @@ def convert_symbol_from_json(
                 },
                 diff_variable=diff_variable,
             )
+        elif symbol_type == "ExpressionFunctionParameter":
+            return ExpressionFunctionParameter(
+                json_data["name"],
+                Serialise.convert_symbol_from_json(json_data["children"][0]),
+                json_data["func_name"],
+                json_data["func_args"],
+            )
         elif symbol_type == "PrimaryBroadcast":
             child = Serialise.convert_symbol_from_json(json_data["children"][0])
             domain = json_data["broadcast_domain"]

src/pybamm/models/full_battery_models/lithium_ion/basic_spm_with_3d_thermal.py

@@ -48,14 +48,14 @@ def __init__(self, options=None, name="SPM with Separate Cell Domain"):
         T = pybamm.Variable("Cell temperature [K]", domain="cell")
 
         if self.options.get("cell geometry") == "pouch":
-            x = pybamm.SpatialVariable("x", domain="cell")
-            y = pybamm.SpatialVariable("y", domain="cell")
-            z = pybamm.SpatialVariable("z", domain="cell")
-            integration_vars = [x, y, z]
+            self.x_cell = pybamm.SpatialVariable("x", domain="cell")
+            self.y_cell = pybamm.SpatialVariable("y", domain="cell")
+            self.z_cell = pybamm.SpatialVariable("z", domain="cell")
+            integration_vars = [self.x_cell, self.y_cell, self.z_cell]
         elif self.options.get("cell geometry") == "cylindrical":
-            r = pybamm.SpatialVariable("r_macro", domain="cell")
-            z = pybamm.SpatialVariable("z", domain="cell")
-            integration_vars = [r, z]
+            self.r_cell = pybamm.SpatialVariable("r_macro", domain="cell")
+            self.z_cell = pybamm.SpatialVariable("z", domain="cell")
+            integration_vars = [self.r_cell, self.z_cell]
         else:
             raise ValueError(
                 f"Geometry type '{self.options.get('cell geometry')}' is not supported. "
@@ -261,9 +261,8 @@ def __init__(self, options=None, name="SPM with Separate Cell Domain"):
     def set_thermal_bcs(self, T):
         geometry_type = self.options.get("cell geometry", "pouch")
         if geometry_type == "pouch":
-            y = pybamm.SpatialVariable("y", "cell")
-            z = pybamm.SpatialVariable("z", "cell")
-            T_amb = self.param.T_amb(y, z, pybamm.t)
+            # Reuse the spatial variables created in __init__
+            T_amb = self.param.T_amb(self.y_cell, self.z_cell, pybamm.t)
             face_params = {
                 "x_min": self.param.h_edge_x_min,
                 "x_max": self.param.h_edge_x_max,
@@ -273,9 +272,8 @@ def set_thermal_bcs(self, T):
                 "z_max": self.param.h_edge_z_max,
             }
         elif geometry_type == "cylindrical":
-            r = pybamm.SpatialVariable("r_macro", "cell")
-            z = pybamm.SpatialVariable("z", "cell")
-            T_amb = self.param.T_amb(r, z, pybamm.t)
+            # Reuse the spatial variables created in __init__
+            T_amb = self.param.T_amb(self.r_cell, self.z_cell, pybamm.t)
             face_params = {
                 "r_min": self.param.h_edge_radial_min,
                 "r_max": self.param.h_edge_radial_max,

src/pybamm/parameters/parameter_values.py

@@ -1,11 +1,14 @@
+import json
 import numbers
 from collections import defaultdict
+from pathlib import Path
 from pprint import pformat
 from warnings import warn
 
 import numpy as np
 
 import pybamm
+from pybamm.expression_tree.operations.serialise import Serialise
 from pybamm.models.full_battery_models.lithium_ion.msmr import (
     is_deprecated_msmr_name,
     replace_deprecated_msmr_name,
@@ -668,7 +671,10 @@ def process_symbol(self, symbol):
         try:
             return self._processed_symbols[symbol]
         except KeyError:
-            processed_symbol = self._process_symbol(symbol)
+            if not isinstance(symbol, pybamm.FunctionParameter):
+                processed_symbol = self._process_symbol(symbol)
+            else:
+                processed_symbol = self._process_function_parameter(symbol)
             self._processed_symbols[symbol] = processed_symbol
 
             return processed_symbol
@@ -835,6 +841,88 @@ def _process_symbol(self, symbol):
             # Backup option: return the object
             return symbol
 
+    def _process_function_parameter(self, symbol):
+        function_parameter = self[symbol.name]
+        # Handle symbolic function parameter case
+        if isinstance(function_parameter, pybamm.ExpressionFunctionParameter):
+            # Process children
+            new_children = []
+            for child in symbol.children:
+                if symbol.diff_variable is not None and any(
+                    x == symbol.diff_variable for x in child.pre_order()
+                ):
+                    # Wrap with NotConstant to avoid simplification,
+                    # which would stop symbolic diff from working properly
+                    new_child = pybamm.NotConstant(child)
+                    new_children.append(self.process_symbol(new_child))
+                else:
+                    new_children.append(self.process_symbol(child))
+
+            # Get the expression and inputs for the function
+            expression = function_parameter.child
+            inputs = {
+                arg: child
+                for arg, child in zip(
+                    function_parameter.func_args, symbol.children, strict=False
+                )
+            }
+
+            # Set domains for function inputs in post-order traversal
+            for node in expression.post_order():
+                if node.name in inputs:
+                    node.domains = inputs[node.name].domains
+                else:
+                    node.domains = node.get_children_domains(node.children)
+
+            # Combine parameter values with inputs
+            combined_params = ParameterValues({**self, **inputs})
+
+            # Process any FunctionParameter children first to avoid recursion
+            for child in expression.pre_order():
+                if isinstance(child, pybamm.FunctionParameter):
+                    # Build new child with parent inputs
+                    new_child_children = [
+                        inputs[child_child.name]
+                        if isinstance(child_child, pybamm.Parameter)
+                        and child_child.name in inputs
+                        else child_child
+                        for child_child in child.children
+                    ]
+                    new_child = pybamm.FunctionParameter(
+                        child.name,
+                        dict(zip(child.input_names, new_child_children, strict=False)),
+                        diff_variable=child.diff_variable,
+                        print_name=child.print_name,
+                    )
+
+                    # For this local combined parameter values, process the new child
+                    # and store the result as the processed symbol for this child
+                    # This means the child is evaluated with the parent inputs only when
+                    # it is called from within the parent function (not elsewhere in
+                    # the expression tree)
+                    combined_params._processed_symbols[child] = (
+                        combined_params.process_symbol(new_child)
+                    )
+
+            # Process function with combined parameter values to get a symbolic
+            # expression
+            function = combined_params.process_symbol(expression)
+
+            # Differentiate if necessary
+            if symbol.diff_variable is None:
+                # Use ones_like so that we get the right shapes
+                function_out = function * pybamm.ones_like(*new_children)
+            else:
+                # return differentiated function
+                new_diff_variable = self.process_symbol(symbol.diff_variable)
+                function_out = function.diff(new_diff_variable)
+
+            return function_out
+
+        # Handle non-symbolic function_name case
+        else:
+            return self._process_symbol(symbol)
+
     def evaluate(self, symbol, inputs=None):
         """
         Process and evaluate a symbol.
@@ -992,3 +1080,32 @@ def __contains__(self, key):
 
     def __iter__(self):
         return iter(self._dict_items)
+
+    @staticmethod
+    def from_json(filename_or_dict):
+        """
+        Loads a ParameterValues object from a JSON file or a dictionary.
+
+        Parameters
+        ----------
+        filename_or_dict : string-like or dict
+            The filename to load the JSON file from, or a dictionary.
+
+        Returns
+        -------
+        ParameterValues
+            The ParameterValues object
+        """
+        if isinstance(filename_or_dict, str | Path):
+            with open(filename_or_dict) as f:
+                parameter_values_dict = json.load(f)
+        elif isinstance(filename_or_dict, dict):
+            parameter_values_dict = filename_or_dict.copy()
+        else:
+            raise TypeError("Input must be a filename (str or pathlib.Path) or a dict")
+
+        for key, value in parameter_values_dict.items():
+            if isinstance(value, dict):
+                parameter_values_dict[key] = Serialise.convert_symbol_from_json(value)
+
+        return ParameterValues(parameter_values_dict)