From 122ca4c6f807e57b7099477ad120c5d30427ff12 Mon Sep 17 00:00:00 2001 From: Robert Timms Date: Mon, 27 Jan 2025 14:18:10 +0000 Subject: [PATCH] #4776 initial fix --- src/pybamm/discretisations/discretisation.py | 7 + src/pybamm/models/base_model.py | 9 +- src/pybamm/solvers/processed_variable.py | 238 +++++++++++++++++- .../solvers/processed_variable_computed.py | 153 ++++++++++- 4 files changed, 400 insertions(+), 7 deletions(-) diff --git a/src/pybamm/discretisations/discretisation.py b/src/pybamm/discretisations/discretisation.py index 1a70f1dd56..5efb67921f 100644 --- a/src/pybamm/discretisations/discretisation.py +++ b/src/pybamm/discretisations/discretisation.py @@ -783,6 +783,13 @@ def process_symbol(self, symbol): ] else: discretised_symbol.secondary_mesh = None + + # Assign tertiary mesh + if symbol.domains["tertiary"] != []: + discretised_symbol.tertiary_mesh = self.mesh[symbol.domains["tertiary"]] + else: + discretised_symbol.tertiary_mesh = None + return discretised_symbol def _process_symbol(self, symbol): diff --git a/src/pybamm/models/base_model.py b/src/pybamm/models/base_model.py index 3c8fa91488..b34e177302 100644 --- a/src/pybamm/models/base_model.py +++ b/src/pybamm/models/base_model.py @@ -122,6 +122,11 @@ def generic_deserialise(cls, instance, properties): else: var.secondary_mesh = None + if var.domains["tertiary"] != []: + var.tertiary_mesh = properties["mesh"][var.domains["tertiary"]] + else: + var.tertiary_mesh = None + if properties["geometry"]: instance._geometry = pybamm.Geometry(properties["geometry"]) else: @@ -875,8 +880,10 @@ def set_initial_conditions_from(self, solution, inplace=True, return_type="model final_state_eval = final_state[:, -1] elif final_state.ndim == 3: final_state_eval = final_state[:, :, -1].flatten(order="F") + elif final_state.ndim == 4: + final_state_eval = final_state[:, :, :, -1].flatten(order="F") else: - raise NotImplementedError("Variable must be 0D, 1D, or 2D") + raise NotImplementedError("Variable must be 0D, 1D, 2D, or 3D") elif isinstance(var, pybamm.Concatenation): children = [] for child in var.orphans: diff --git a/src/pybamm/solvers/processed_variable.py b/src/pybamm/solvers/processed_variable.py index 30fbcedc70..1e49426734 100644 --- a/src/pybamm/solvers/processed_variable.py +++ b/src/pybamm/solvers/processed_variable.py @@ -860,6 +860,221 @@ def _interp_setup(self, entries, t): return entries, coords_for_interp +class ProcessedVariable3D(ProcessedVariable): + """ + An object that can be evaluated at arbitrary (scalars or vectors) t and x, and + returns the (interpolated) value of the base variable at that t and x. + + Parameters + ---------- + base_variables : list of :class:`pybamm.Symbol` + A list of base variables with a method `evaluate(t,y)`, each entry of which + returns the value of that variable for that particular sub-solution. + A Solution can be comprised of sub-solutions which are the solutions of + different models. + Note that this can be any kind of node in the expression tree, not + just a :class:`pybamm.Variable`. + When evaluated, returns an array of size (m,n) + base_variables_casadi : list of :class:`casadi.Function` + A list of casadi functions. When evaluated, returns the same thing as + `base_Variable.evaluate` (but more efficiently). + solution : :class:`pybamm.Solution` + The solution object to be used to create the processed variables + """ + + def __init__( + self, + base_variables, + base_variables_casadi, + solution, + time_integral: Optional[pybamm.ProcessedVariableTimeIntegral] = None, + ): + self.dimensions = 3 + super().__init__( + base_variables, + base_variables_casadi, + solution, + time_integral=time_integral, + ) + first_dim_nodes = self.mesh.nodes + first_dim_edges = self.mesh.edges + second_dim_nodes = self.base_variables[0].secondary_mesh.nodes + third_dim_nodes = self.base_variables[0].tertiary_mesh.nodes + if self.base_eval_size // (len(second_dim_nodes) * len(third_dim_nodes)) == len( + first_dim_nodes + ): + first_dim_pts = first_dim_nodes + elif self.base_eval_size // ( + len(second_dim_nodes) * len(third_dim_nodes) + ) == len(first_dim_edges): + first_dim_pts = first_dim_edges + + second_dim_pts = second_dim_nodes + third_dim_pts = third_dim_nodes + self.first_dim_size = len(first_dim_pts) + self.second_dim_size = len(second_dim_pts) + self.third_dim_size = len(third_dim_pts) + + def _observe_raw_python(self): + """ + Initialise a 3D object that depends on x, y, and z or x, r, and R. + """ + pybamm.logger.debug("Observing the variable raw data in Python") + first_dim_size, second_dim_size, t_size = self._shape(self.t_pts) + entries = np.empty((first_dim_size, second_dim_size, t_size)) + + # Evaluate the base_variable index-by-index + idx = 0 + for ts, ys, inputs, base_var_casadi in zip( + self.all_ts, self.all_ys, self.all_inputs_casadi, self.base_variables_casadi + ): + for inner_idx, t in enumerate(ts): + t = ts[inner_idx] + y = ys[:, inner_idx] + entries[:, :, idx] = np.reshape( + base_var_casadi(t, y, inputs).full(), + [first_dim_size, second_dim_size], + order="F", + ) + idx += 1 + return entries + + def _interp_setup(self, entries, t): + """ + Initialise a 3D object that depends on x, y, and z, or x, r, and R. + """ + first_dim_nodes = self.mesh.nodes + first_dim_edges = self.mesh.edges + second_dim_nodes = self.base_variables[0].secondary_mesh.nodes + second_dim_edges = self.base_variables[0].secondary_mesh.edges + third_dim_nodes = self.base_variables[0].tertiary_mesh.nodes + third_dim_edges = self.base_variables[0].tertiary_mesh.edges + if self.base_eval_size // (len(second_dim_nodes) * len(third_dim_nodes)) == len( + first_dim_nodes + ): + first_dim_pts = first_dim_nodes + elif self.base_eval_size // ( + len(second_dim_nodes) * len(third_dim_nodes) + ) == len(first_dim_edges): + first_dim_pts = first_dim_edges + + second_dim_pts = second_dim_nodes + third_dim_pts = third_dim_nodes + + # add points outside first dimension domain for extrapolation to + # boundaries + extrap_space_first_dim_left = np.array( + [2 * first_dim_pts[0] - first_dim_pts[1]] + ) + extrap_space_first_dim_right = np.array( + [2 * first_dim_pts[-1] - first_dim_pts[-2]] + ) + first_dim_pts = np.concatenate( + [extrap_space_first_dim_left, first_dim_pts, extrap_space_first_dim_right] + ) + extrap_entries_left = np.expand_dims(2 * entries[0] - entries[1], axis=0) + extrap_entries_right = np.expand_dims(2 * entries[-1] - entries[-2], axis=0) + entries_for_interp = np.concatenate( + [extrap_entries_left, entries, extrap_entries_right], axis=0 + ) + + # add points outside second dimension domain for extrapolation to + # boundaries + extrap_space_second_dim_left = np.array( + [2 * second_dim_pts[0] - second_dim_pts[1]] + ) + extrap_space_second_dim_right = np.array( + [2 * second_dim_pts[-1] - second_dim_pts[-2]] + ) + second_dim_pts = np.concatenate( + [ + extrap_space_second_dim_left, + second_dim_pts, + extrap_space_second_dim_right, + ] + ) + extrap_entries_second_dim_left = np.expand_dims( + 2 * entries_for_interp[:, 0, :] - entries_for_interp[:, 1, :], axis=1 + ) + extrap_entries_second_dim_right = np.expand_dims( + 2 * entries_for_interp[:, -1, :] - entries_for_interp[:, -2, :], axis=1 + ) + entries_for_interp = np.concatenate( + [ + extrap_entries_second_dim_left, + entries_for_interp, + extrap_entries_second_dim_right, + ], + axis=1, + ) + + # add points outside tertiary dimension domain for extrapolation to + # boundaries + extrap_space_third_dim_left = np.array( + [2 * third_dim_pts[0] - third_dim_pts[1]] + ) + extrap_space_third_dim_right = np.array( + [2 * third_dim_pts[-1] - third_dim_pts[-2]] + ) + third_dim_pts = np.concatenate( + [ + extrap_space_third_dim_left, + third_dim_pts, + extrap_space_third_dim_right, + ] + ) + extrap_entries_third_dim_left = np.expand_dims( + 2 * entries_for_interp[:, :, 0] - entries_for_interp[:, :, 1], axis=2 + ) + extrap_entries_third_dim_right = np.expand_dims( + 2 * entries_for_interp[:, :, -1] - entries_for_interp[:, :, -2], axis=2 + ) + entries_for_interp = np.concatenate( + [ + extrap_entries_third_dim_left, + entries_for_interp, + extrap_entries_third_dim_right, + ], + axis=2, + ) + + self.spatial_variable_names = { + k: self._process_spatial_variable_names(v) + for k, v in self.spatial_variables.items() + } + + self.first_dimension = self.spatial_variable_names["primary"] + self.second_dimension = self.spatial_variable_names["secondary"] + self.third_dimension = self.spatial_variable_names["tertiary"] + + # assign attributes for reference + first_dim_pts_for_interp = first_dim_pts + second_dim_pts_for_interp = second_dim_pts + third_dim_pts_for_interp = third_dim_pts + + # Set pts to edges for nicer plotting + self.first_dim_pts = first_dim_edges + self.second_dim_pts = second_dim_edges + self.third_dim_pts = third_dim_edges + + # save attributes for interpolation + coords_for_interp = { + self.first_dimension: first_dim_pts_for_interp, + self.second_dimension: second_dim_pts_for_interp, + self.third_dimension: third_dim_pts_for_interp, + "t": t, + } + + return entries_for_interp, coords_for_interp + + def _shape(self, t): + first_dim_size = self.first_dim_size + second_dim_size = self.second_dim_size + third_dim_size = self.third_dim_size + t_size = len(t) + return [first_dim_size, second_dim_size, third_dim_size, t_size] + + def process_variable(base_variables, *args, **kwargs): mesh = base_variables[0].mesh domain = base_variables[0].domain @@ -875,6 +1090,15 @@ def process_variable(base_variables, *args, **kwargs): and isinstance(mesh, pybamm.ScikitSubMesh2D) ): return ProcessedVariable2DSciKitFEM(base_variables, *args, **kwargs) + if ( + base_variables[0].secondary_mesh + and "current collector" in base_variables[0].domains["secondary"] + and isinstance(base_variables[0].secondary_mesh, pybamm.ScikitSubMesh2D) + ): + raise NotImplementedError( + "3D variables with secondary domain 'current collector' using the ScikitFEM" + " discretisation are not supported as processed variables" + ) # check variable shape if len(base_eval_shape) == 0 or base_eval_shape[0] == 1: @@ -896,11 +1120,15 @@ def process_variable(base_variables, *args, **kwargs): ]: return ProcessedVariable2D(base_variables, *args, **kwargs) - # Raise error for 3D variable - raise NotImplementedError( - f"Shape not recognized for {base_variables[0]}" - + "(note processing of 3D variables is not yet implemented)" - ) + # Try some shapes that could make the variable a 3D variable + tertiary_pts = base_variables[0].tertiary_mesh.nodes + if base_eval_size // (len(second_dim_pts) * len(tertiary_pts)) in [ + len(first_dim_nodes), + len(first_dim_edges), + ]: + return ProcessedVariable3D(base_variables, *args, **kwargs) + + raise NotImplementedError(f"Shape not recognized for {base_variables[0]}") def _is_f_contiguous(all_ys): diff --git a/src/pybamm/solvers/processed_variable_computed.py b/src/pybamm/solvers/processed_variable_computed.py index befe6314b6..90be679ca3 100644 --- a/src/pybamm/solvers/processed_variable_computed.py +++ b/src/pybamm/solvers/processed_variable_computed.py @@ -16,7 +16,7 @@ class ProcessedVariableComputed: The 'Computed' variant of ProcessedVariable deals with variables that have been derived at solve time (see the 'output_variables' solver option), - where the full state-vector is not itself propogated and returned. + where the full state-vector is not itself propagated and returned. Parameters ---------- @@ -423,6 +423,157 @@ def initialise_2D_scikit_fem(self): coords={"y": y_sol, "z": z_sol, "t": self.t_pts}, ) + def initialise_3D(self): + """ + Initialise a 3D object that depends on x, y, and z, or x, r, and R. + """ + first_dim_nodes = self.mesh.nodes + first_dim_edges = self.mesh.edges + second_dim_nodes = self.base_variables[0].secondary_mesh.nodes + second_dim_edges = self.base_variables[0].secondary_mesh.edges + third_dim_nodes = self.base_variables[0].tertiary_mesh.nodes + third_dim_edges = self.base_variables[0].tertiary_mesh.edges + if self.base_eval_size // (len(second_dim_nodes) * len(third_dim_nodes)) == len( + first_dim_nodes + ): + first_dim_pts = first_dim_nodes + elif self.base_eval_size // ( + len(second_dim_nodes) * len(third_dim_nodes) + ) == len(first_dim_edges): + first_dim_pts = first_dim_edges + + second_dim_pts = second_dim_nodes + third_dim_pts = third_dim_nodes + + first_dim_size = len(first_dim_pts) + second_dim_size = len(second_dim_pts) + third_dim_size = len(third_dim_pts) + + entries = self.unroll_3D( + realdata=None, + n_dim1=third_dim_size, + n_dim2=second_dim_size, + n_dim3=first_dim_size, + axis_swaps=[(0, 3), (0, 2), (0, 1)], + ) + + # add points outside first dimension domain for extrapolation to + # boundaries + extrap_space_first_dim_left = np.array( + [2 * first_dim_pts[0] - first_dim_pts[1]] + ) + extrap_space_first_dim_right = np.array( + [2 * first_dim_pts[-1] - first_dim_pts[-2]] + ) + first_dim_pts = np.concatenate( + [extrap_space_first_dim_left, first_dim_pts, extrap_space_first_dim_right] + ) + extrap_entries_left = np.expand_dims(2 * entries[0] - entries[1], axis=0) + extrap_entries_right = np.expand_dims(2 * entries[-1] - entries[-2], axis=0) + entries_for_interp = np.concatenate( + [extrap_entries_left, entries, extrap_entries_right], axis=0 + ) + + # add points outside second dimension domain for extrapolation to + # boundaries + extrap_space_second_dim_left = np.array( + [2 * second_dim_pts[0] - second_dim_pts[1]] + ) + extrap_space_second_dim_right = np.array( + [2 * second_dim_pts[-1] - second_dim_pts[-2]] + ) + second_dim_pts = np.concatenate( + [ + extrap_space_second_dim_left, + second_dim_pts, + extrap_space_second_dim_right, + ] + ) + extrap_entries_second_dim_left = np.expand_dims( + 2 * entries_for_interp[:, 0, :] - entries_for_interp[:, 1, :], axis=1 + ) + extrap_entries_second_dim_right = np.expand_dims( + 2 * entries_for_interp[:, -1, :] - entries_for_interp[:, -2, :], axis=1 + ) + entries_for_interp = np.concatenate( + [ + extrap_entries_second_dim_left, + entries_for_interp, + extrap_entries_second_dim_right, + ], + axis=1, + ) + + # add points outside tertiary dimension domain for extrapolation to + # boundaries + extrap_space_third_dim_left = np.array( + [2 * third_dim_pts[0] - third_dim_pts[1]] + ) + extrap_space_third_dim_right = np.array( + [2 * third_dim_pts[-1] - third_dim_pts[-2]] + ) + third_dim_pts = np.concatenate( + [ + extrap_space_third_dim_left, + third_dim_pts, + extrap_space_third_dim_right, + ] + ) + extrap_entries_third_dim_left = np.expand_dims( + 2 * entries_for_interp[:, :, 0] - entries_for_interp[:, :, 1], axis=2 + ) + extrap_entries_third_dim_right = np.expand_dims( + 2 * entries_for_interp[:, :, -1] - entries_for_interp[:, :, -2], axis=2 + ) + entries_for_interp = np.concatenate( + [ + extrap_entries_third_dim_left, + entries_for_interp, + extrap_entries_third_dim_right, + ], + axis=2, + ) + + # Process r-x, x-z, r-R, R-x, or R-z + if ( + self.domain[0].endswith("particle") + and self.domains["secondary"][0].endswith("particle size") + and self.domains["tertiary"][0].endswith("electrode") + ): + self.first_dimension = "r" + self.second_dimension = "R" + self.third_dimension = "x" + self.r_sol = first_dim_pts + self.R_sol = second_dim_pts + self.x_sol = third_dim_pts + else: # pragma: no cover + raise pybamm.DomainError( + f"Cannot process 3D object with domains '{self.domains}'." + ) + + # assign attributes for reference + self.entries = entries + self.dimensions = 3 + first_dim_pts_for_interp = first_dim_pts + second_dim_pts_for_interp = second_dim_pts + third_dim_pts_for_interp = third_dim_pts + + # Set pts to edges for nicer plotting + self.first_dim_pts = first_dim_edges + self.second_dim_pts = second_dim_edges + self.third_dim_pts = third_dim_edges + + # set up interpolation + self._xr_data_array = xr.DataArray( + entries_for_interp, + coords={ + self.first_dimension: first_dim_pts_for_interp, + self.second_dimension: second_dim_pts_for_interp, + self.third_dimension: third_dim_pts_for_interp, + "t": self.t_pts, + }, + ) + def __call__(self, t=None, x=None, r=None, y=None, z=None, R=None): """ Evaluate the variable at arbitrary *dimensional* t (and x, r, y, z and/or R),