#1100 merge 1221

Author: valentinsulzer

examples/scripts/DFN.py

@@ -30,7 +30,7 @@
 
 # solve model
 t_eval = np.linspace(0, 3600, 100)
-solver = pybamm.CasadiSolver(mode="fast", atol=1e-6, rtol=1e-3)
+solver = pybamm.CasadiSolver(mode="safe", atol=1e-6, rtol=1e-3)
 solution = solver.solve(model, t_eval)
 
 # plot

examples/scripts/experimental_protocols/cccv.py

@@ -4,18 +4,18 @@
 import pybamm
 import matplotlib.pyplot as plt
 
-pybamm.set_logging_level("INFO")
+pybamm.set_logging_level("DEBUG")
 experiment = pybamm.Experiment(
     [
-        "Discharge at C/10 for 10 hours or until 3.3 V",
+        "Discharge at C/1 for 1 hours or until 3.3 V",
         "Rest for 1 hour",
         "Charge at 1 A until 4.1 V",
         "Hold at 4.1 V until 50 mA",
         "Rest for 1 hour",
     ]
     * 3
 )
-model = pybamm.lithium_ion.DFN()
+model = pybamm.lithium_ion.SPM()
 sim = pybamm.Simulation(model, experiment=experiment, solver=pybamm.CasadiSolver())
 sim.solve()
 

pybamm/models/base_model.py

@@ -109,6 +109,7 @@ def __init__(self, name="Unnamed model"):
         self.external_variables = []
         self._parameters = None
         self._input_parameters = None
+        self._variables_casadi = {}
 
         # Default behaviour is to use the jacobian and simplify
         self.use_jacobian = True
@@ -117,6 +118,7 @@ def __init__(self, name="Unnamed model"):
 
         # Model is not initially discretised
         self.is_discretised = False
+        self.y_slices = None
 
         # Default timescale is 1 second
         self.timescale = pybamm.Scalar(1)
@@ -340,6 +342,14 @@ def new_empty_copy(self):
         new_model.convert_to_format = self.convert_to_format
         new_model.timescale = self.timescale
         new_model.length_scales = self.length_scales
+
+        # Variables from discretisation
+        new_model.is_discretised = self.is_discretised
+        new_model.y_slices = self.y_slices
+        new_model.concatenated_rhs = self.concatenated_rhs
+        new_model.concatenated_algebraic = self.concatenated_algebraic
+        new_model.concatenated_initial_conditions = self.concatenated_initial_conditions
+
         return new_model
 
     def new_copy(self):
@@ -425,6 +435,31 @@ def set_initial_conditions_from(self, solution, inplace=True):
                     "Variable must have type 'Variable' or 'Concatenation'"
                 )
 
+        # Also update the concatenated initial conditions if the model is already
+        # discretised
+        if model.is_discretised:
+            # Unpack slices for sorting
+            y_slices = {var.id: slce for var, slce in model.y_slices.items()}
+            slices = []
+            for symbol in model.initial_conditions.keys():
+                if isinstance(symbol, pybamm.Concatenation):
+                    # must append the slice for the whole concatenation, so that
+                    # equations get sorted correctly
+                    slices.append(
+                        slice(
+                            y_slices[symbol.children[0].id][0].start,
+                            y_slices[symbol.children[-1].id][0].stop,
+                        )
+                    )
+                else:
+                    slices.append(y_slices[symbol.id][0])
+            equations = list(model.initial_conditions.values())
+            # sort equations according to slices
+            sorted_equations = [eq for _, eq in sorted(zip(slices, equations))]
+            model.concatenated_initial_conditions = pybamm.NumpyConcatenation(
+                *sorted_equations
+            )
+
         return model
 
     def check_and_combine_dict(self, dict1, dict2):
@@ -901,10 +936,7 @@ def default_spatial_methods(self):
     @property
     def default_solver(self):
         "Return default solver based on whether model is ODE model or DAE model"
-        if len(self.algebraic) == 0:
-            return pybamm.ScipySolver()
-        else:
-            return pybamm.CasadiSolver(mode="safe")
+        return pybamm.CasadiSolver(mode="safe")
 
 
 # helper functions for finding symbols

pybamm/simulation.py

@@ -532,17 +532,16 @@ def get_variable_array(self, *variables):
             arrays.
         """
 
-        variable_arrays = [
-            self.built_model.variables[var].evaluate(
-                self.solution.t[-1], self.solution.y[:, -1]
-            )
-            for var in variables
-        ]
-
-        if len(variable_arrays) == 1:
-            return variable_arrays[0]
-        else:
-            return tuple(variable_arrays)
+        variable_arrays = {}
+        for var in variables:
+            processed_var = self.solution[var].data
+            if processed_var.ndim == 1:
+                variable_arrays[var] = processed_var[-1]
+            elif processed_var.ndim == 2:
+                variable_arrays[var] = processed_var[:, -1]
+            elif processed_var.ndim == 3:
+                variable_arrays[var] = processed_var[:, :, -1]
+        return variable_arrays
 
     def plot(self, output_variables=None, quick_plot_vars=None, **kwargs):
         """
@@ -693,6 +692,8 @@ def save(self, filename):
             and self._solver.integrator_specs != {}
         ):
             self._solver.integrator_specs = {}
+        if self.solution is not None:
+            self.solution.clear_casadi_attributes()
         with open(filename, "wb") as f:
             pickle.dump(self, f, pickle.HIGHEST_PROTOCOL)
 

pybamm/solvers/base_solver.py

@@ -1058,7 +1058,6 @@ def __init__(self, function, name, model):
         self.timescale = self.model.timescale_eval
 
     def __call__(self, t, y, inputs):
-        y = y.reshape(-1, 1)
         if self.name in ["rhs", "algebraic", "residuals"]:
             pybamm.logger.debug(
                 "Evaluating {} for {} at t={}".format(

pybamm/solvers/casadi_solver.py

@@ -155,8 +155,6 @@ def _integrate(self, model, t_eval, inputs=None):
             return solution
         elif self.mode in ["safe", "safe without grid"]:
             y0 = model.y0
-            if isinstance(y0, casadi.DM):
-                y0 = y0.full().flatten()
             # Step-and-check
             t = t_eval[0]
             t_f = t_eval[-1]
@@ -174,7 +172,7 @@ def _integrate(self, model, t_eval, inputs=None):
                 self.create_integrator(model, inputs_dict)
                 # Initialize solution
                 solution = pybamm.Solution(
-                    np.array([t]), y0[:, np.newaxis], model=model, inputs=inputs_dict
+                    np.array([t]), y0, model=model, inputs=inputs_dict
                 )
                 solution.solve_time = 0
                 solution.integration_time = 0
@@ -457,7 +455,7 @@ def _run_integrator(self, model, y0, inputs_dict, t_eval):
                     x0=y0_diff, z0=y0_alg, p=inputs_eval, **self.extra_options_call
                 )
                 integration_time = timer.time()
-                y_sol = np.concatenate([sol["xf"].full(), sol["zf"].full()])
+                y_sol = casadi.vertcat(sol["xf"], sol["zf"])
                 sol = pybamm.Solution(t_eval, y_sol)
                 sol.integration_time = integration_time
                 return sol

pybamm/solvers/dummy_solver.py

@@ -34,7 +34,7 @@ def _integrate(self, model, t_eval, inputs=None):
         """
         y_sol = np.zeros((1, t_eval.size))
         sol = pybamm.Solution(
-            t_eval, y_sol, termination="final time", model=model, inputs=inputs_dict
+            t_eval, y_sol, termination="final time", model=model, inputs=inputs
         )
         sol.integration_time = 0
         return sol

pybamm/solvers/jax_solver.py

@@ -203,7 +203,7 @@ def _integrate(self, model, t_eval, inputs=None):
         t_event = None
         y_event = onp.array(None)
         sol = pybamm.Solution(
-            t_eval, y, t_event, y_event, termination, model=model, inputs=inputs_dict
+            t_eval, y, t_event, y_event, termination, model=model, inputs=inputs
         )
         sol.integration_time = integration_time
         return sol

pybamm/solvers/processed_variable.py

@@ -40,24 +40,25 @@ class ProcessedVariable(object):
         variable. 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_variable_casadi : :class:`casadi.Function`
+        A casadi function. 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
-    known_evals : dict
-        Dictionary of known evaluations, to be used to speed up finding the solution
     warn : bool, optional
         Whether to raise warnings when trying to evaluate time and length scales.
         Default is True.
     """
 
-    def __init__(self, base_variable, solution, known_evals=None, warn=True):
+    def __init__(self, base_variable, base_variable_casadi, solution, warn=True):
         self.base_variable = base_variable
+        self.base_variable_casadi = base_variable_casadi
         self.t_sol = solution.t
         self.u_sol = solution.y
         self.mesh = base_variable.mesh
         self.inputs = solution.inputs
         self.domain = base_variable.domain
         self.auxiliary_domains = base_variable.auxiliary_domains
-        self.known_evals = known_evals
         self.warn = warn
 
         # Sensitivity starts off uninitialized, only set when called
@@ -104,19 +105,10 @@ def __init__(self, base_variable, solution, known_evals=None, warn=True):
             self.length_scales = solution.length_scales_eval
 
         # Evaluate base variable at initial time
-        if self.known_evals:
-            self.base_eval, self.known_evals[solution.t[0]] = base_variable.evaluate(
-                self.t_sol[0],
-                self.u_sol[:, 0],
-                inputs={name: inp[:, 0] for name, inp in solution.inputs.items()},
-                known_evals=self.known_evals[solution.t[0]],
-            )
-        else:
-            self.base_eval = base_variable.evaluate(
-                solution.t[0],
-                solution.y[:, 0],
-                inputs={name: inp[:, 0] for name, inp in solution.inputs.items()},
-            )
+        inputs = casadi.vertcat(*[inp[:, 0] for inp in self.inputs.values()])
+        self.base_eval = self.base_variable_casadi(
+            solution.t[0], solution.y[:, 0], inputs
+        ).full()
 
         # handle 2D (in space) finite element variables differently
         if (
@@ -164,13 +156,8 @@ def initialise_0D(self):
         for idx in range(len(self.t_sol)):
             t = self.t_sol[idx]
             u = self.u_sol[:, idx]
-            inputs = {name: inp[:, idx] for name, inp in self.inputs.items()}
-            if self.known_evals:
-                entries[idx], self.known_evals[t] = self.base_variable.evaluate(
-                    t, u, inputs=inputs, known_evals=self.known_evals[t]
-                )
-            else:
-                entries[idx] = self.base_variable.evaluate(t, u, inputs=inputs)
+            inputs = casadi.vertcat(*[inp[:, idx] for inp in self.inputs.values()])
+            entries[idx] = self.base_variable_casadi(t, u, inputs).full()[0, 0]
 
         # set up interpolation
         if len(self.t_sol) == 1:
@@ -200,15 +187,8 @@ def initialise_1D(self, fixed_t=False):
         for idx in range(len(self.t_sol)):
             t = self.t_sol[idx]
             u = self.u_sol[:, idx]
-            inputs = {name: inp[:, idx] for name, inp in self.inputs.items()}
-            if self.known_evals:
-                eval_and_known_evals = self.base_variable.evaluate(
-                    t, u, inputs=inputs, known_evals=self.known_evals[t]
-                )
-                entries[:, idx] = eval_and_known_evals[0][:, 0]
-                self.known_evals[t] = eval_and_known_evals[1]
-            else:
-                entries[:, idx] = self.base_variable.evaluate(t, u, inputs=inputs)[:, 0]
+            inputs = casadi.vertcat(*[inp[:, idx] for inp in self.inputs.values()])
+            entries[:, idx] = self.base_variable_casadi(t, u, inputs).full()[:, 0]
 
         # Get node and edge values
         nodes = self.mesh.nodes
@@ -310,23 +290,12 @@ def initialise_2D(self):
         for idx in range(len(self.t_sol)):
             t = self.t_sol[idx]
             u = self.u_sol[:, idx]
-            inputs = {name: inp[:, idx] for name, inp in self.inputs.items()}
-            if self.known_evals:
-                eval_and_known_evals = self.base_variable.evaluate(
-                    t, u, inputs=inputs, known_evals=self.known_evals[t]
-                )
-                entries[:, :, idx] = np.reshape(
-                    eval_and_known_evals[0],
-                    [first_dim_size, second_dim_size],
-                    order="F",
-                )
-                self.known_evals[t] = eval_and_known_evals[1]
-            else:
-                entries[:, :, idx] = np.reshape(
-                    self.base_variable.evaluate(t, u, inputs=inputs),
-                    [first_dim_size, second_dim_size],
-                    order="F",
-                )
+            inputs = casadi.vertcat(*[inp[:, idx] for inp in self.inputs.values()])
+            entries[:, :, idx] = np.reshape(
+                self.base_variable_casadi(t, u, inputs).full(),
+                [first_dim_size, second_dim_size],
+                order="F",
+            )
 
         # add points outside first dimension domain for extrapolation to
         # boundaries
@@ -463,22 +432,13 @@ def initialise_2D_scikit_fem(self):
         for idx in range(len(self.t_sol)):
             t = self.t_sol[idx]
             u = self.u_sol[:, idx]
-            inputs = {name: inp[:, idx] for name, inp in self.inputs.items()}
+            inputs = casadi.vertcat(*[inp[:, idx] for inp in self.inputs.values()])
 
-            if self.known_evals:
-                eval_and_known_evals = self.base_variable.evaluate(
-                    t, u, inputs=inputs, known_evals=self.known_evals[t]
-                )
-                entries[:, :, idx] = np.reshape(
-                    eval_and_known_evals[0], [len_y, len_z], order="F"
-                )
-                self.known_evals[t] = eval_and_known_evals[1]
-            else:
-                entries[:, :, idx] = np.reshape(
-                    self.base_variable.evaluate(t, u, inputs=inputs),
-                    [len_y, len_z],
-                    order="F",
-                )
+            entries[:, :, idx] = np.reshape(
+                self.base_variable_casadi(t, u, inputs).full(),
+                [len_y, len_z],
+                order="F",
+            )
 
         # assign attributes for reference
         self.entries = entries