diff --git a/CHANGELOG.md b/CHANGELOG.md index 49d8030374..888aad896c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,7 @@ - Improve the performance of matrix multiplication with CasADi expressions. ([#5351](https://github.com/pybamm-team/PyBaMM/pull/5351)) - Adds option for lists of inputs to `solve` to include input parameters which are used as initial conditions. ([#5311](https://github.com/pybamm-team/PyBaMM/pull/5311)) +- DiffSL export for discretised PyBaMM models. ([#5370](https://github.com/pybamm-team/PyBaMM/pull/5370)) ## Bug fixes diff --git a/pyproject.toml b/pyproject.toml index aedb0e7c18..229a28809a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -110,6 +110,8 @@ dev = [ "importlib-metadata; python_version < '3.10'", # For property based testing "hypothesis", + # For diffsl tests (not available on Mac Intel x86_64) + "pydiffsol; platform_system != 'Darwin' or platform_machine != 'x86_64'", ] [project.entry-points."pybamm_parameter_sets"] diff --git a/src/pybamm/__init__.py b/src/pybamm/__init__.py index 734a710aca..1258626580 100644 --- a/src/pybamm/__init__.py +++ b/src/pybamm/__init__.py @@ -67,6 +67,9 @@ from .models.event import Event from .models.event import EventType +# DiffSL export +from .expression_tree.operations.diffsl import DiffSLExport + # Battery models from .models.full_battery_models.base_battery_model import ( BaseBatteryModel, diff --git a/src/pybamm/expression_tree/operations/diffsl.py b/src/pybamm/expression_tree/operations/diffsl.py new file mode 100644 index 0000000000..0ec03420fc --- /dev/null +++ b/src/pybamm/expression_tree/operations/diffsl.py @@ -0,0 +1,697 @@ +from itertools import chain + +import numpy as np +from scipy.sparse import csr_matrix + +import pybamm + + +class DiffSLExport: + """ + A class to export parameterised and discretised PyBaMM models to the DiffSL format. + + Attributes: + model (pybamm.BaseModel): The PyBaMM model to be exported. This model should be + parameterised and discretised before exporting. + float_precision (int): The number of significant digits for float representation. + """ + + def __init__(self, model: pybamm.BaseModel, float_precision: int = 20): + """ + Initializes the DiffSLExport class. + + Args: + model (pybamm.BaseModel): The model to export. + float_precision (int): The number of significant digits for float representation. + + Raises: + ValueError: If float_precision is not a positive integer. + """ + self.model = model + if not isinstance(float_precision, int) or float_precision <= 0: + raise ValueError("float_precision must be a positive integer") + self.float_precision = float_precision + + @staticmethod + def _name_tensor( + symbol: pybamm.Symbol, + tensor_index: int, + is_variable: bool = False, + is_event: bool = False, + ) -> str: + if is_variable: + tensor_name = f"variable{tensor_index}" + elif is_event: + tensor_name = f"event{tensor_index}" + else: + tensor_name = f"varying{tensor_index}" + return tensor_name + + @staticmethod + def scalar_to_diffeq( + symbol: pybamm.Scalar, tensor_index: int, float_precision: int = 20 + ) -> tuple[str, str]: + tensor_name = f"constant{tensor_index}" + tensor_def = ( + f"{tensor_name}_i " + + "{" + + f"\n (0:1): {symbol.value:.{float_precision}g},\n" + + "}" + ) + return tensor_name, tensor_def + + @staticmethod + def vector_to_diffeq( + symbol: pybamm.Vector, tensor_index: int, float_precision: int = 20 + ) -> tuple[str, str]: + if isinstance(symbol.entries, csr_matrix): + vector = symbol.entries.toarray().flatten() + elif isinstance(symbol.entries, np.ndarray): + vector = symbol.entries.flatten() + else: + raise TypeError( + f"{type(symbol.entries)} not implemented" + ) # pragma: no cover + + tensor_name = f"constant{tensor_index}" + + # any segments of the vector that are the same value are combined + curr_value = vector[0] + start_index = 0 + lines = [f"{tensor_name}_i " + "{"] + for i in range(1, vector.size): + if vector[i] != curr_value: + end_index = i + lines += [ + f" ({start_index}:{end_index}): {curr_value:.{float_precision}g}," + ] + start_index = i + curr_value = vector[i] + end_index = vector.size + lines += [f" ({start_index}:{end_index}): {curr_value:.{float_precision}g},"] + new_line = "\n" + return tensor_name, new_line.join(lines) + new_line + "}" + + @staticmethod + def matrix_to_diffeq( + symbol: pybamm.Matrix, tensor_index: int, float_precision: int = 20 + ) -> tuple[str, str]: + tensor_name = f"constant{tensor_index}" + if isinstance(symbol.entries, csr_matrix): + nrows, ncols = symbol.entries.shape + lines = [f"{tensor_name}_ij " + "{"] + is_diagonal = nrows == ncols + is_constant = True + is_row_vector = nrows == 1 + max_colj = 0 + max_rowi = 0 + for rowi in range(nrows): + for j in range( + symbol.entries.indptr[rowi], symbol.entries.indptr[rowi + 1] + ): + colj = symbol.entries.indices[j] + value = symbol.entries.data[j] + max_colj = max(max_colj, colj) + max_rowi = max(max_rowi, rowi) + if abs(value - symbol.entries.data[0]) > 1e-12: + is_constant = False + if rowi != colj: + is_diagonal = False + if is_diagonal and is_constant and symbol.entries.nnz == min(nrows, ncols): + min_dim = min(nrows, ncols) + lines += [ + f" (0..{min_dim - 1},0..{min_dim - 1}): {symbol.entries.data[0]:.{float_precision}g}," + ] + elif is_row_vector and is_constant and symbol.entries.nnz == ncols: + lines += [ + f" (0,0:{ncols}): {symbol.entries.data[0]:.{float_precision}g}," + ] + else: + for rowi in range(nrows): + for j in reversed( + range( + symbol.entries.indptr[rowi], symbol.entries.indptr[rowi + 1] + ) + ): + colj = symbol.entries.indices[j] + lines += [ + f" ({rowi},{colj}): {symbol.entries.data[j]:.{float_precision}g}," + ] + + if max_colj < ncols - 1 or max_rowi < nrows - 1: + # add a zero entry to the end to make sure the matrix is the right size + lines += [f" ({nrows - 1},{ncols - 1}): 0.0,"] + + elif isinstance(symbol.entries, np.ndarray): + nrows, ncols = symbol.entries.shape + lines = [f"{tensor_name}_ij " + "{"] + if nrows == 1: + # any segments of the row vector that are the same value are combined + curr_value = symbol.entries[0, 0] + start_index = 0 + for coli in range(ncols): + value = symbol.entries[0, coli] + if value != curr_value: + end_index = coli + lines += [ + f" (0, {start_index}:{end_index}): {curr_value:.{float_precision}g}," + ] + start_index = coli + curr_value = value + end_index = ncols + lines += [ + f" (0, {start_index}:{end_index}): {curr_value:.{float_precision}g}," + ] + else: + for rowi in range(nrows): + for colj in range(ncols): + lines += [ + f" ({rowi},{colj}): {symbol.entries[rowi, colj]:.{float_precision}g}," + ] + else: + raise TypeError( + f"{type(symbol.entries)} not implemented" + ) # pragma: no cover + new_line = "\n" + return tensor_name, new_line.join(lines) + new_line + "}" + + def extract_pre_calculated_tensors( + self, + eqn: pybamm.Symbol, + symbol_to_tensor_name: dict[pybamm.Symbol, str], + tensor_index: int, + y_slice_to_label: dict[tuple[int], str], + diffeq: dict[str, str], + symbol_counts: dict[pybamm.Symbol, int], + is_variable: bool = False, + is_event: bool = False, + ) -> int: + new_line = "\n" + for symbol in eqn.post_order(): + # extract any binary operators that occur more than twice and dont involve scalars + if ( + isinstance( + symbol, + pybamm.BinaryOperator | pybamm.UnaryOperator | pybamm.Function, + ) + and symbol_counts.get(symbol, 0) > 1 + ): + if symbol in symbol_to_tensor_name: + continue + has_scalar = any( + [isinstance(child, pybamm.Scalar) for child in symbol.children] + ) + if has_scalar: + continue + tensor_name = DiffSLExport._name_tensor( + symbol, tensor_index, is_variable, is_event + ) + tensor_index += 1 + lines = [f"{tensor_name}_i " + "{"] + eqn_str = equation_to_diffeq( + symbol, + y_slice_to_label, + symbol_to_tensor_name, + float_precision=self.float_precision, + ) + lines += [f" {eqn_str},"] + symbol_to_tensor_name[symbol] = tensor_name + diffeq[tensor_name] = new_line.join(lines) + new_line + "}" + # skip top-level + for child in eqn.children: + for symbol in child.post_order(): + if ( + isinstance(symbol, pybamm.BinaryOperator) + and symbol.name == "@" + and isinstance(symbol.left, pybamm.Matrix) + ): + # extract the matrix vector product as a separate tensor + if symbol in symbol_to_tensor_name: + continue + tensor_name = DiffSLExport._name_tensor( + symbol, tensor_index, is_variable, is_event + ) + tensor_index += 1 + lines = [f"{tensor_name}_i " + "{"] + eqn = equation_to_diffeq( + symbol, + y_slice_to_label, + symbol_to_tensor_name, + float_precision=self.float_precision, + ) + lines += [f" {eqn},"] + symbol_to_tensor_name[symbol] = tensor_name + diffeq[tensor_name] = new_line.join(lines) + new_line + "}" + + elif isinstance(symbol, pybamm.DomainConcatenation): + if symbol in symbol_to_tensor_name: + continue + tensor_name = DiffSLExport._name_tensor( + symbol, tensor_index, is_variable, is_event + ) + tensor_index += 1 + lines = [f"{tensor_name}_i " + "{"] + for child, slices in zip( + symbol.children, symbol._children_slices, strict=False + ): + eqn = equation_to_diffeq( + child, + y_slice_to_label, + symbol_to_tensor_name, + float_precision=self.float_precision, + ) + for child_dom, child_slice in slices.items(): + for i, _slice in enumerate(child_slice): + s = symbol._slices[child_dom][i] + lines += [f" ({s.start}:{s.stop}): {eqn},"] + symbol_to_tensor_name[symbol] = tensor_name + diffeq[tensor_name] = new_line.join(lines) + new_line + "}" + + return tensor_index + + def to_diffeq(self, outputs: list[str]) -> str: + """Convert a pybamm model to a diffeq model""" + model = self.model.new_copy() + all_vars = model.get_processed_variables_dict() + if len(all_vars) == 0 and len(model.variables) > 0: + all_vars = model.variables + if not isinstance(outputs, list) or any( + not isinstance(o, str) for o in outputs + ): + raise TypeError("outputs must be a list of str") + if len(outputs) == 0: + raise ValueError("outputs must be a non-empty list of str") + for out in outputs: + if out not in all_vars: + print("all_vars keys:", all_vars.keys()) # DEBUG + raise ValueError(f"output {out} not in model") + eqn = all_vars[out] + has_events = len(model.events) > 0 + is_ode = len(model.algebraic) == 0 + + states = list(chain(model.rhs.keys(), model.algebraic.keys())) + state_labels = [to_variable_name(v.name) for v in states] + initial_conditions = [model.initial_conditions[v] for v in states] + diffeq = {} + new_line = "\n" + + # inputs, find all pybamm.InputParameters in the model + inputs = [] + for eqn in chain( + model.rhs.values(), + model.algebraic.values(), + [all_vars[output] for output in outputs], + [e.expression for e in model.events], + ): + for symbol in eqn.pre_order(): + if isinstance(symbol, pybamm.InputParameter): + input_name = to_variable_name(symbol.name) + if input_name not in inputs: + inputs.append(input_name) + + if len(inputs) > 0: + lines = ["in_i {"] + for inpt in inputs: + lines += [f" {to_variable_name(inpt)} = 1,"] + diffeq["in"] = new_line.join(lines) + new_line + "}" + + # extract constant vectors and matrices from model as tensors + symbol_to_tensor_name = {} + vectors: set[pybamm.Vector] = set() + matrices: set[pybamm.Matrix] = set() + termination_events = [ + e.expression + for e in model.events + if e.event_type == pybamm.EventType.TERMINATION + ] + for eqn in chain( + model.rhs.values(), + model.algebraic.values(), + [all_vars[output] for output in outputs], + termination_events, + ): + for symbol in eqn.pre_order(): + if isinstance(symbol, pybamm.Vector): + vectors.add(symbol) + elif isinstance(symbol, pybamm.Matrix): + matrices.add(symbol) + + tensor_index = 0 + for symbol in vectors: + tensor_name, tensor_def = DiffSLExport.vector_to_diffeq( + symbol, tensor_index, float_precision=self.float_precision + ) + tensor_index += 1 + symbol_to_tensor_name[symbol] = tensor_name + diffeq[tensor_name] = tensor_def + + for symbol in matrices: + tensor_name, tensor_def = DiffSLExport.matrix_to_diffeq( + symbol, tensor_index, float_precision=self.float_precision + ) + tensor_index += 1 + symbol_to_tensor_name[symbol] = tensor_name + diffeq[tensor_name] = tensor_def + + # state vector u + input_lines = [] + lines = ["u_i {"] + start_index = 0 + y_slice_to_label = {} + new_line = "\n" + for i, ic in enumerate(initial_conditions): + if isinstance(ic, pybamm.Vector): + tensor_name, tensor_def = DiffSLExport.vector_to_diffeq( + ic, tensor_index, float_precision=self.float_precision + ) + elif isinstance(ic, pybamm.Scalar): + tensor_name, tensor_def = DiffSLExport.scalar_to_diffeq( + ic, tensor_index, float_precision=self.float_precision + ) + else: + raise TypeError( + f"Initial condition of type {type(ic)} not supported" + ) # pragma: no cover + input_lines += [tensor_def] + tensor_index += 1 + label = state_labels[i] + y_slice_to_label[(start_index, start_index + ic.size)] = label + lines += [f" {label} = {tensor_name}_i,"] + start_index += ic.size + diffeq["u"] = ( + new_line.join(input_lines) + + new_line + + new_line.join(lines) + + new_line + + "}" + ) + + # diff of state vector u + if not is_ode: + lines = ["dudt_i {"] + start_index = 0 + for i, ic in enumerate(initial_conditions): + label = state_name_to_dstate_name(state_labels[i]) + indices = ( + f"({start_index}:{start_index + ic.size}): " if ic.size > 1 else "" + ) + zero = pybamm.Scalar(0) + eqn = equation_to_diffeq( + zero, + {}, + symbol_to_tensor_name, + float_precision=self.float_precision, + ) + lines += [f" {indices}{label} = {eqn},"] + start_index += ic.size + diffeq["dudt"] = new_line.join(lines) + new_line + "}" + + symbol_counts = {} + for eqn in chain( + model.rhs.values(), + model.algebraic.values(), + termination_events, + [all_vars[output] for output in outputs], + ): + for s in eqn.pre_order(): + if s in symbol_counts: + symbol_counts[s] += 1 + else: + symbol_counts[s] = 1 + + for eqn in chain( + model.rhs.values(), + model.algebraic.values(), + ): + tensor_index = self.extract_pre_calculated_tensors( + eqn, + symbol_to_tensor_name, + tensor_index, + y_slice_to_label, + diffeq, + symbol_counts, + ) + + for eqn in chain( + termination_events, + ): + tensor_index = self.extract_pre_calculated_tensors( + eqn, + symbol_to_tensor_name, + tensor_index, + y_slice_to_label, + diffeq, + symbol_counts, + is_event=True, + ) + + for eqn in chain( + [all_vars[output] for output in outputs], + ): + tensor_index = self.extract_pre_calculated_tensors( + eqn, + symbol_to_tensor_name, + tensor_index, + y_slice_to_label, + diffeq, + symbol_counts, + is_variable=True, + ) + + # M + if not is_ode: + lines = ["M_i {"] + start_index = 0 + for i, rhs in enumerate(model.rhs.values()): + eqn = f"{state_name_to_dstate_name(state_labels[i])}_i" + lines += [f" {eqn},"] + start_index += rhs.size + for algebraic in model.algebraic.values(): + n = algebraic.size + if n == 1: + eqn = "0.0" + else: + eqn = f"({start_index}:{start_index + n}): 0.0" + lines += [f" {eqn},"] + start_index += algebraic.size + diffeq["M"] = new_line.join(lines) + new_line + "}" + + # F + lines = ["F_i {"] + for rhs in model.rhs.values(): + eqn = equation_to_diffeq( + rhs, + y_slice_to_label, + symbol_to_tensor_name, + float_precision=self.float_precision, + ) + lines += [f" {eqn},"] + start_index += rhs.size + for algebraic in model.algebraic.values(): + eqn = equation_to_diffeq( + algebraic, + y_slice_to_label, + symbol_to_tensor_name, + float_precision=self.float_precision, + ) + lines += [f" {eqn},"] + start_index += algebraic.size + diffeq["F"] = new_line.join(lines) + new_line + "}" + + # out + lines = ["out_i {"] + for output in outputs: + eqn = equation_to_diffeq( + all_vars[output], + y_slice_to_label, + symbol_to_tensor_name, + float_precision=self.float_precision, + ) + lines += [f" {eqn},"] + diffeq["out"] = new_line.join(lines) + new_line + "}" + + # events + if has_events: + lines = ["stop_i {"] + for event in model.events: + if event.event_type == pybamm.EventType.TERMINATION: + eqn = equation_to_diffeq( + event.expression, + y_slice_to_label, + symbol_to_tensor_name, + float_precision=self.float_precision, + ) + lines += [f" {eqn},"] + diffeq["stop"] = new_line.join(lines) + new_line + "}" + + all_lines = [] + + if is_ode: + state_tensors = ["u"] + f_and_g = ["F"] + else: + state_tensors = ["u", "dudt"] + f_and_g = ["M", "F"] + if has_events: + stop = ["stop"] + else: + stop = [] + out = ["out"] + + # inputs and constants + if "in" in diffeq: + all_lines = [diffeq["in"]] + for key in diffeq.keys(): + if key.startswith("constant"): + all_lines += [diffeq[key]] + for key in state_tensors: + all_lines += [diffeq[key]] + for key in diffeq.keys(): + if key.startswith("varying"): + all_lines += [diffeq[key]] + for key in f_and_g: + all_lines += [diffeq[key]] + for key in diffeq.keys(): + if key.startswith("event"): + all_lines += [diffeq[key]] + for key in stop: + all_lines += [diffeq[key]] + for key in diffeq.keys(): + if key.startswith("variable"): + all_lines += [diffeq[key]] + for key in out: + all_lines += [diffeq[key]] + + return "\n".join(all_lines) + + +def _equation_to_diffeq( + equation: pybamm.Symbol, + y_slice_to_label: dict[tuple[int], str], + symbol_to_tensor_name: dict[pybamm.Symbol, str], + float_precision: int = 20, + transpose: bool = False, +) -> str: + if equation in symbol_to_tensor_name: + if isinstance(equation, pybamm.Matrix): + index = "ji" if transpose else "ij" + else: + index = "j" if transpose else "i" + return f"{symbol_to_tensor_name[equation]}_{index}" + if isinstance(equation, pybamm.BinaryOperator): + left = _equation_to_diffeq( + equation.left, + y_slice_to_label, + symbol_to_tensor_name, + float_precision=float_precision, + transpose=transpose, + ) + + if equation.name == "@" and isinstance(equation.left, pybamm.Matrix): + right = _equation_to_diffeq( + equation.right, + y_slice_to_label, + symbol_to_tensor_name, + float_precision=float_precision, + transpose=not transpose, + ) + return f"({left} * {right})" + else: + right = _equation_to_diffeq( + equation.right, + y_slice_to_label, + symbol_to_tensor_name, + float_precision=float_precision, + transpose=transpose, + ) + + if equation.name == "maximum": + return f"max({left}, {right})" + elif equation.name == "minimum": + return f"min({left}, {right})" + elif equation.name == "@": + return f"({left} * {right})" + elif equation.name == "**": + return f"pow({left}, {right})" + return f"({left} {equation.name} {right})" + elif isinstance(equation, pybamm.UnaryOperator): + return f"{equation.name}({_equation_to_diffeq(equation.child, y_slice_to_label, symbol_to_tensor_name, float_precision=float_precision, transpose=transpose)})" + elif isinstance(equation, pybamm.Function): + name = equation.function.__name__ + args = ", ".join( + [ + _equation_to_diffeq( + x, + y_slice_to_label, + symbol_to_tensor_name, + float_precision=float_precision, + transpose=transpose, + ) + for x in equation.children + ] + ) + return f"{name}({args})" + elif isinstance(equation, pybamm.Scalar): + return f"{equation.value:.{float_precision}g}" + elif isinstance(equation, pybamm.StateVector): + if len(equation.y_slices) != 1: + raise NotImplementedError( + "only one state vector slice supported" + ) # pragma: no cover + start = equation.y_slices[0].start + end = equation.y_slices[0].stop + index = "j" if transpose else "i" + # look for exact match + for sl, label in y_slice_to_label.items(): + if sl[0] == start and sl[1] == end: + return f"{label}_{index}" + # didn't find exact match, look for containing slice + for sl, label in y_slice_to_label.items(): + if sl[0] <= start and sl[1] >= end: + start_offset = start - sl[0] + end_offset = end - sl[0] + return f"{label}_{index}[{start_offset}:{end_offset}]" + raise ValueError( + f"equation {equation} not in slice_to_label" + ) # pragma: no cover + + elif isinstance(equation, pybamm.InputParameter): + return f"{to_variable_name(equation.name)}" + elif isinstance(equation, pybamm.Time): + return "t" + else: + raise TypeError( + f"{type(equation)} not implemented for symbol {equation}" + ) # pragma: no cover + + +def equation_to_diffeq( + equation: pybamm.Symbol, + slice_to_label: dict[tuple[int], str], + symbol_to_tensor_name: dict[pybamm.Symbol, str], + float_precision: int = 20, + transpose: bool = False, +) -> str: + if not isinstance(equation, pybamm.Symbol): + raise TypeError("equation must be a pybamm.Symbol") # pragma: no cover + return _equation_to_diffeq( + equation, + slice_to_label, + symbol_to_tensor_name, + float_precision=float_precision, + transpose=transpose, + ) + + +def to_variable_name(name: str) -> str: + """Convert a name to a valid diffeq variable name""" + convert_to_underscore = [" ", "-", "(", ")", "[", "]", "{", "}", "/", "\\", "."] + name = name.lower() + for char in convert_to_underscore: + name = name.replace(char, "") + return name + + +def state_name_to_dstate_name(state_name: str) -> str: + """Convert a state name to a dstate name""" + return f"d{state_name}dt" diff --git a/tests/integration/test_solvers/test_diffsl_export.py b/tests/integration/test_solvers/test_diffsl_export.py new file mode 100644 index 0000000000..83ff7e5c55 --- /dev/null +++ b/tests/integration/test_solvers/test_diffsl_export.py @@ -0,0 +1,119 @@ +# +# Tests for DiffSLExport class +# + +import importlib.util +import logging +import time + +import numpy as np +import pytest + +import pybamm + +has_pydiffsol = importlib.util.find_spec("pydiffsol") is not None + + +class TestDiffSLExport: + @pytest.mark.skipif(not has_pydiffsol, reason="pydiffsol is not installed") + @pytest.mark.parametrize( + "model", + [ + pybamm.lithium_ion.SPM(), + pybamm.lithium_ion.SPMe(), + pybamm.lithium_ion.DFN(), + ], + ids=["SPM", "SPMe", "DFN"], + ) + @pytest.mark.parametrize( + "inputs", + [[], ["Current function [A]"], ["Ambient temperature [K]"]], + ids=["no_inputs", "with_current", "with_temperature"], + ) + def test_models(self, model, inputs): + import pydiffsol as ds + + pv = model.default_parameter_values + pv_inputs = {} + ds_inputs = [] + for input_name in inputs: + pv_inputs[input_name] = 0.9 * pv[input_name] + ds_inputs.append(0.9 * pv[input_name]) + pv[input_name] = "[input]" + output_variable = "Voltage [V]" + + t0 = time.perf_counter() + geometry = model.default_geometry + model_param = pv.process_model(model, inplace=False) + pv.process_geometry(geometry) + spatial_methods = model.default_spatial_methods + var_pts = model.default_var_pts + submesh_types = model.default_submesh_types + mesh = pybamm.Mesh(geometry, submesh_types, var_pts) + disc = pybamm.Discretisation(mesh, spatial_methods) + model_disc = disc.process_model(model_param, inplace=True) + logger = logging.getLogger() + logger.info( + f"Pybamm simulation creation time: {time.perf_counter() - t0:.5f} seconds" + ) + + t0 = time.perf_counter() + diffsl_code = pybamm.DiffSLExport(model_disc).to_diffeq( + outputs=[output_variable] + ) + logger.info(f"DiffSL export time: {time.perf_counter() - t0:.5f} seconds") + + t0 = time.perf_counter() + ode = ds.Ode( + diffsl_code, + matrix_type=ds.faer_sparse, + scalar_type=ds.f64, + linear_solver=ds.lu, + method=ds.tr_bdf2, + ) + logger.info( + f"DiffSL Ode compilation time: {time.perf_counter() - t0:.5f} seconds" + ) + ode.rtol = 1e-4 + ode.atol = 1e-6 + if isinstance(model, pybamm.lithium_ion.DFN): + ode.ic_options.armijo_constant = 1e-1 + + t_eval = [0, 3600] + t_interp = np.linspace(t_eval[0], t_eval[1], 100) + + solver = pybamm.IDAKLUSolver() + soln_pybamm = solver.solve( + model_disc, t_eval, t_interp=t_interp, inputs=pv_inputs + ).y + + t0 = time.perf_counter() + voltage_pybamm = solver.solve( + model_disc, t_eval, t_interp=t_interp, inputs=pv_inputs + )[output_variable].data + logger.info(f"Pybamm solve time: {time.perf_counter() - t0:.5f} seconds") + + n = model_disc.y0.shape[0] + v = np.ones(n) + for i in range(soln_pybamm.shape[1]): + y = soln_pybamm[:, i] + pybamm_jac = model_disc.jac_rhs_algebraic_action_eval( + 0, y, np.array(ds_inputs), v + ) + pybamm_dydt = model_disc.rhs_algebraic_eval(0, y, np.array(ds_inputs)) + diffsol_dydt = ode.rhs(np.array(ds_inputs), 0, y.flatten()) + diffsol_jac = ode.rhs_jac_mul(np.array(ds_inputs), 0, y.flatten(), v) + np.testing.assert_allclose( + pybamm_dydt.full().flatten(), diffsol_dydt, rtol=1e-5, atol=1e-8 + ) + np.testing.assert_allclose( + pybamm_jac.full().flatten(), diffsol_jac, rtol=1e-5, atol=1e-8 + ) + + ds_inputs = np.array(ds_inputs) + t0 = time.perf_counter() + voltage_diffsol = ode.solve_dense(np.array(ds_inputs), t_interp) + logger.info(f"DiffSL solve time: {time.perf_counter() - t0:.5f} seconds") + if voltage_diffsol.shape[0] == 1: + voltage_diffsol = voltage_diffsol.flatten() + np.testing.assert_allclose(voltage_pybamm, voltage_diffsol, rtol=3e-4) diff --git a/tests/unit/test_expression_tree/test_operations/test_diffsl_export.py b/tests/unit/test_expression_tree/test_operations/test_diffsl_export.py new file mode 100644 index 0000000000..694217d59e --- /dev/null +++ b/tests/unit/test_expression_tree/test_operations/test_diffsl_export.py @@ -0,0 +1,159 @@ +# +# Tests for DiffSLExport class +# +import numpy as np +import pytest +import scipy.sparse + +import pybamm + + +class TestDiffSLExport: + # fixture of a model with all the characteristics we want to test + @pytest.fixture + def model(self): + model = pybamm.BaseModel() + + x = pybamm.StateVector(slice(0, 2), domain="negative electrode") + y = pybamm.StateVector(slice(2, 4), domain="positive electrode") + z = pybamm.StateVector(slice(4, 5)) + z2 = pybamm.StateVector(slice(5, 7)) + A = pybamm.Matrix( + scipy.sparse.csr_matrix(([4.12345, 4.12345], [0, 1], [0, 1, 2])) + ) + B = pybamm.Matrix(np.array([[4, -2], [3, -1]])) + C = pybamm.Matrix(scipy.sparse.csr_matrix(([1], [0], [0, 1, 1]), shape=(2, 2))) + b = pybamm.Matrix(scipy.sparse.csr_matrix(([2, 2], [0, 1], [0, 2]))) + c = pybamm.Vector(scipy.sparse.csr_matrix(([1], [0], [0, 1, 1]))) + d = pybamm.Matrix(np.array([2, 3]).reshape((1, 2))) + u = pybamm.StateVector(slice(0, 4)) + p = pybamm.InputParameter("p") + dup = pybamm.maximum(x * x, 0) + dxdt = A @ dup + c + b @ dup + C @ dup + d @ x + A @ pybamm.minimum(x, 0) + dydt = A @ y + y @ z + pybamm.cos(p**2) + pybamm.t + + x_n = pybamm.SpatialVariable("x_n", domain="negative electrode") + x_p = pybamm.SpatialVariable("x_p", domain="positive electrode") + geometry = pybamm.Geometry( + { + "negative electrode": { + x_n: {"min": pybamm.Scalar(0), "max": pybamm.Scalar(1)} + }, + "positive electrode": { + x_p: {"min": pybamm.Scalar(1), "max": pybamm.Scalar(2)} + }, + } + ) + submesh_types = { + "negative electrode": pybamm.Uniform1DSubMesh, + "positive electrode": pybamm.Uniform1DSubMesh, + } + var_pts = {x_n: 2, x_p: 2} + mesh = pybamm.Mesh(geometry, submesh_types, var_pts) + dudt = pybamm.DomainConcatenation([dxdt, dydt], mesh) + + model.rhs = {u: u * dudt + u / dudt} + model.algebraic = {z: z, z2: z2} + model.initial_conditions = { + u: pybamm.Vector(np.array([1, 2, 1, 2])), + z: pybamm.Scalar(0), + z2: pybamm.Vector(np.array([0, 0])), + } + model.variables = {"x": x, "y": y, "z": z} + model.events = [pybamm.Event("event1", x - B @ x)] + return model + + def test_model(self, model): + export = pybamm.DiffSLExport(model, float_precision=6).to_diffeq(outputs=["x"]) + assert "u_i" in export + assert "event" in export + assert "constant" in export + assert "varying" in export + assert "dudt_i" in export + assert "F_i" in export + assert "out_i" in export + assert "stop_i" in export + assert "in_i" in export + assert "cos" in export + assert "max" in export + assert "min" in export + + def test_float_precision(self, model): + export = pybamm.DiffSLExport(model, float_precision=6).to_diffeq(outputs=["x"]) + assert "4.12345" in export + export = pybamm.DiffSLExport(model, float_precision=2).to_diffeq(outputs=["x"]) + assert "4.12345" not in export + assert "4.1" in export + with pytest.raises(ValueError): + model = pybamm.DiffSLExport(model, float_precision=-1) + + def test_inputs(self, model): + with pytest.raises(TypeError): + pybamm.DiffSLExport(model).to_diffeq(outputs="not a list") + with pytest.raises(ValueError): + pybamm.DiffSLExport(model).to_diffeq(outputs=["not in model"]) + with pytest.raises(ValueError): + pybamm.DiffSLExport(model).to_diffeq(outputs=[]) + + def test_ode(self): + model = pybamm.BaseModel() + + x = pybamm.Variable("x") + y = pybamm.Variable("y") + + dxdt = 4 * x - 2 * y + dydt = 3 * x - y + + model.rhs = {x: dxdt, y: dydt} + model.initial_conditions = {x: pybamm.Scalar(1), y: pybamm.Scalar(2)} + model.variables = {"x": x, "y": y, "z": x + 4 * y} + + disc = pybamm.Discretisation() + model = disc.process_model(model) + + model = pybamm.DiffSLExport(model) + correct_export = "constant0_i {\n (0:1): 1,\n}\nconstant1_i {\n (0:1): 2,\n}\nu_i {\n x = constant0_i,\n y = constant1_i,\n}\nF_i {\n ((4 * x_i) - (2 * y_i)),\n ((3 * x_i) - y_i),\n}\nout_i {\n x_i,\n y_i,\n (x_i + (4 * y_i)),\n}" + assert correct_export == model.to_diffeq(outputs=["x", "y", "z"]) + + def test_heat_equation(self): + model = pybamm.BaseModel() + + x = pybamm.SpatialVariable("x", domain="rod", coord_sys="cartesian") + T = pybamm.Variable("Temperature", domain="rod") + k = pybamm.Parameter("Thermal diffusivity") + + N = -k * pybamm.grad(T) + dTdt = -pybamm.div(N) + model.rhs = {T: dTdt} + + model.boundary_conditions = { + T: { + "left": (pybamm.Scalar(0), "Dirichlet"), + "right": (pybamm.Scalar(0), "Dirichlet"), + } + } + + model.initial_conditions = {T: x} + model.variables = {"Temperature": T, "Heat flux": N} + geometry = {"rod": {x: {"min": pybamm.Scalar(0), "max": pybamm.Scalar(10)}}} + param = pybamm.ParameterValues({"Thermal diffusivity": 1.0}) + param.process_model(model) + param.process_geometry(geometry) + + submesh_types = {"rod": pybamm.Uniform1DSubMesh} + var_pts = {x: 10} + mesh = pybamm.Mesh(geometry, submesh_types, var_pts) + spatial_methods = {"rod": pybamm.FiniteVolume()} + disc = pybamm.Discretisation(mesh, spatial_methods) + + disc.process_model(model) + + model = pybamm.DiffSLExport(model) + + # The order of the matrix entries is not deterministic, so we need to check both possible options + correct_export_option_1 = "constant0_ij {\n (0,0): 2,\n (1,0): -1,\n (1,1): 1,\n (2,1): -1,\n (2,2): 1,\n (3,2): -1,\n (3,3): 1,\n (4,3): -1,\n (4,4): 1,\n (5,4): -1,\n (5,5): 1,\n (6,5): -1,\n (6,6): 1,\n (7,6): -1,\n (7,7): 1,\n (8,7): -1,\n (8,8): 1,\n (9,8): -1,\n (9,9): 1,\n (10,9): -2,\n}\nconstant1_ij {\n (0,0): -3,\n (0,1): 1,\n (1,1): -2,\n (1,0): 1,\n (1,2): 1,\n (2,2): -2,\n (2,1): 1,\n (2,3): 1,\n (3,3): -2,\n (3,2): 1,\n (3,4): 1,\n (4,4): -2,\n (4,3): 1,\n (4,5): 1,\n (5,5): -2,\n (5,4): 1,\n (5,6): 1,\n (6,6): -2,\n (6,5): 1,\n (6,7): 1,\n (7,7): -2,\n (7,6): 1,\n (7,8): 1,\n (8,8): -2,\n (8,7): 1,\n (8,9): 1,\n (9,9): -3,\n (9,8): 1,\n}\nconstant2_i {\n (0:1): 0.5,\n (1:2): 1.5,\n (2:3): 2.5,\n (3:4): 3.5,\n (4:5): 4.5,\n (5:6): 5.5,\n (6:7): 6.5,\n (7:8): 7.5,\n (8:9): 8.5,\n (9:10): 9.5,\n}\nu_i {\n temperature = constant2_i,\n}\nF_i {\n (constant1_ij * temperature_j),\n}\nvariable3_i {\n (constant0_ij * temperature_j),\n}\nout_i {\n -(variable3_i),\n temperature_i,\n}" + correct_export_option_2 = "constant0_ij {\n (0,0): -3,\n (0,1): 1,\n (1,1): -2,\n (1,0): 1,\n (1,2): 1,\n (2,2): -2,\n (2,1): 1,\n (2,3): 1,\n (3,3): -2,\n (3,2): 1,\n (3,4): 1,\n (4,4): -2,\n (4,3): 1,\n (4,5): 1,\n (5,5): -2,\n (5,4): 1,\n (5,6): 1,\n (6,6): -2,\n (6,5): 1,\n (6,7): 1,\n (7,7): -2,\n (7,6): 1,\n (7,8): 1,\n (8,8): -2,\n (8,7): 1,\n (8,9): 1,\n (9,9): -3,\n (9,8): 1,\n}\nconstant1_ij {\n (0,0): 2,\n (1,0): -1,\n (1,1): 1,\n (2,1): -1,\n (2,2): 1,\n (3,2): -1,\n (3,3): 1,\n (4,3): -1,\n (4,4): 1,\n (5,4): -1,\n (5,5): 1,\n (6,5): -1,\n (6,6): 1,\n (7,6): -1,\n (7,7): 1,\n (8,7): -1,\n (8,8): 1,\n (9,8): -1,\n (9,9): 1,\n (10,9): -2,\n}\nconstant2_i {\n (0:1): 0.5,\n (1:2): 1.5,\n (2:3): 2.5,\n (3:4): 3.5,\n (4:5): 4.5,\n (5:6): 5.5,\n (6:7): 6.5,\n (7:8): 7.5,\n (8:9): 8.5,\n (9:10): 9.5,\n}\nu_i {\n temperature = constant2_i,\n}\nF_i {\n (constant0_ij * temperature_j),\n}\nvariable3_i {\n (constant1_ij * temperature_j),\n}\nout_i {\n -(variable3_i),\n temperature_i,\n}" + model_output = model.to_diffeq(outputs=["Heat flux", "Temperature"]) + assert correct_export_option_1 == str( + model_output + ) or correct_export_option_2 == str(model_output)