#1996 working pass at nomemalloc

Author: valentinsulzer

pybamm/solvers/base_solver.py

@@ -1294,7 +1294,7 @@ def step(
             t = old_solution.all_ts[-1][-1]
             if old_solution.all_models[-1] == model:
                 # initialize with old solution
-                model.y0 = old_solution.all_ys[-1][:, -1]
+                model.y0 = old_solution.y_last
             else:
                 model.y0 = model.set_initial_conditions_from(
                     old_solution

pybamm/solvers/casadi_solver.py

@@ -284,7 +284,7 @@ def _integrate(self, model, t_eval, inputs_dict=None):
                     # update time
                     t = t_window[-1]
                     # update y0
-                    y0 = solution.all_ys[-1][:, -1]
+                    y0 = solution.y_last
 
             # now we extract sensitivities from the solution
             if bool(model.calculate_sensitivities):
@@ -312,7 +312,7 @@ def find_t_event(sol, typ):
 
             # Check most recent y to see if any events have been crossed
             if model.terminate_events_eval:
-                y_last = sol.all_ys[-1][:, -1]
+                y_last = sol.y_last
                 crossed_events = np.sign(
                     init_event_signs
                     * np.concatenate(
@@ -676,7 +676,8 @@ def _run_integrator(
         else:
             integrator = self.integrators[model]["no grid"]
 
-        len_rhs = model.concatenated_rhs.size
+        len_rhs = model.len_rhs
+        len_t = len(t_eval)
 
         # Check y0 to see if it includes sensitivities
         if explicit_sensitivities:
@@ -699,17 +700,20 @@ def _run_integrator(
                     x0=y0_diff, z0=y0_alg, p=inputs_with_tmin, **self.extra_options_call
                 )
                 integration_time = timer.time()
+                x_sols = casadi.horzsplit(casadi_sol["xf"])
                 if casadi_sol["zf"].is_empty():
-                    y_sol = casadi_sol["xf"]
+                    y_sols = x_sols
                 else:
-                    y_sol = pybamm.NoMemAllocVertcat(casadi_sol["xf"], casadi_sol["zf"])
+                    z_sols = casadi.horzsplit(casadi_sol["zf"])
+                    y_sols = [
+                        pybamm.NoMemAllocVertcat(x, z) for x, z in zip(x_sols, z_sols)
+                    ]
             else:
                 # Repeated calls to the integrator
+                y_sols = [y0]
                 x = y0_diff
                 z = y0_alg
-                y_diff = x
-                y_alg = z
-                for i in range(len(t_eval) - 1):
+                for i in range(len_t - 1):
                     t_min = t_eval[i]
                     t_max = t_eval[i + 1]
                     inputs_with_tlims = casadi.vertcat(inputs, t_min, t_max)
@@ -720,19 +724,16 @@ def _run_integrator(
                     integration_time = timer.time()
                     x = casadi_sol["xf"]
                     z = casadi_sol["zf"]
-                    y_diff = casadi.horzcat(y_diff, x)
-                    if not z.is_empty():
-                        y_alg = casadi.horzcat(y_alg, z)
-                if z.is_empty():
-                    y_sol = y_diff
-                else:
-                    y_sol = pybamm.NoMemAllocVertcat(y_diff, y_alg)
+                    if z.is_empty():
+                        y_sols.append(x)
+                    else:
+                        y_sols.append(pybamm.NoMemAllocVertcat(x, z))
 
             sol = pybamm.Solution(
-                t_eval,
-                y_sol,
-                model,
-                inputs_dict,
+                np.array_split(t_eval, len_t),
+                y_sols,
+                [model] * len_t,
+                [inputs_dict] * len_t,
                 sensitivities=extract_sensitivities_in_solution,
                 check_solution=False,
             )

pybamm/solvers/processed_variable.py

@@ -39,9 +39,14 @@ def __init__(self, base_variables, base_variables_casadi, solution, warn=True):
 
         self.all_ts = solution.all_ts
         self.all_ys = solution.all_ys
+        self.y_last = solution.y_last
         self.all_inputs = solution.all_inputs
         self.all_inputs_casadi = solution.all_inputs_casadi
 
+        # utilities for no_mem_alloc
+        self.uses_no_mem_alloc = isinstance(self.all_ys[0], pybamm.NoMemAllocVertcat)
+        self.out = casadi.DM.zeros(self.all_ys[0].shape)
+
         self.mesh = base_variables[0].mesh
         self.domain = base_variables[0].domain
         self.domains = base_variables[0].domains
@@ -60,9 +65,9 @@ def __init__(self, base_variables, base_variables_casadi, solution, warn=True):
         # Store length scales
         self.length_scales = solution.length_scales_eval
 
-        # Evaluate base variable at initial time
+        # Evaluate base variable at final time
         self.base_eval = self.base_variables_casadi[0](
-            self.all_ts[0][0], self.all_ys[0][:, 0], self.all_inputs_casadi[0]
+            self.all_ts[-1][-1], self.y_last, self.all_inputs_casadi[-1]
         ).full()
 
         # handle 2D (in space) finite element variables differently
@@ -114,7 +119,7 @@ def initialise_0D(self):
         ):
             for inner_idx, t in enumerate(ts):
                 t = ts[inner_idx]
-                y = ys[:, inner_idx]
+                y = self.getitem(ys, idx, self.out)
                 entries[idx] = base_var_casadi(t, y, inputs).full()[0, 0]
                 idx += 1
 
@@ -146,7 +151,7 @@ def initialise_1D(self, fixed_t=False):
         ):
             for inner_idx, t in enumerate(ts):
                 t = ts[inner_idx]
-                y = ys[:, inner_idx]
+                y = self.getitem(ys, idx, self.out)
                 entries[:, idx] = base_var_casadi(t, y, inputs).full()[:, 0]
                 idx += 1
 
@@ -247,7 +252,7 @@ def initialise_2D(self):
         ):
             for inner_idx, t in enumerate(ts):
                 t = ts[inner_idx]
-                y = ys[:, inner_idx]
+                y = self.getitem(ys, idx, self.out)
                 entries[:, :, idx] = np.reshape(
                     base_var_casadi(t, y, inputs).full(),
                     [first_dim_size, second_dim_size],
@@ -419,7 +424,7 @@ def initialise_2D_scikit_fem(self):
         ):
             for inner_idx, t in enumerate(ts):
                 t = ts[inner_idx]
-                y = ys[:, inner_idx]
+                y = self.getitem(ys, idx, self.out)
                 entries[:, :, idx] = np.reshape(
                     base_var_casadi(t, y, inputs).full(),
                     [len_y, len_z],
@@ -454,6 +459,13 @@ def initialise_2D_scikit_fem(self):
                 bounds_error=False,
             )
 
+    def getitem(self, ys, idx, out):
+        if self.uses_no_mem_alloc:
+            ys.get_value(out)
+            return out
+        else:
+            return ys[:, idx]
+
     def __call__(self, t=None, x=None, r=None, y=None, z=None, R=None, warn=True):
         """
         Evaluate the variable at arbitrary *dimensional* t (and x, r, y, z and/or R),

pybamm/solvers/solution.py

@@ -81,6 +81,7 @@ def __init__(
             all_ys = [all_ys]
         if not isinstance(all_models, list):
             all_models = [all_models]
+
         self._all_ts = all_ts
         self._all_ys = all_ys
         self._all_ys_and_sens = all_ys
@@ -96,6 +97,25 @@ def __init__(
         else:
             self.all_inputs = all_inputs
 
+        if not (
+            len(self.all_ts) == len(self.all_ys)
+            or len(self.all_ts) == 1
+            or len(self.all_ys) == 1
+        ):
+            raise ValueError("all_ts and all_ys must be the same length")
+        if not (
+            len(self.all_ts) == len(self.all_models)
+            or len(self.all_ts) == 1
+            or len(self.all_models) == 1
+        ):
+            raise ValueError("all_ts and all_models must be the same length")
+        if not (
+            len(self.all_ts) == len(self.all_inputs)
+            or len(self.all_ts) == 1
+            or len(self.all_inputs) == 1
+        ):
+            raise ValueError("all_ts and all_inputs must be the same length")
+
         self.sensitivities = sensitivities
 
         self._t_event = t_event
@@ -106,7 +126,6 @@ def __init__(
             isinstance(v, casadi.MX) for v in self.all_inputs[0].values()
         )
 
-        # Check no ys are too large
         if check_solution and not self.has_symbolic_inputs:
             self.check_ys_are_not_too_large()
 
@@ -321,8 +340,7 @@ def check_ys_are_not_too_large(self):
         # Only check last one so that it doesn't take too long
         # We only care about the cases where y is growing too large without any
         # restraint, so if y gets large in the middle then comes back down that is ok
-        y, model = self.all_ys[-1], self.all_models[-1]
-        y = y[:, -1]
+        y, model = self.y_last, self.all_models[-1]
         if np.any(y > pybamm.settings.max_y_value):
             for var in [*model.rhs.keys(), *model.algebraic.keys()]:
                 y_var = y[model.variables[var.name].y_slices[0]]
@@ -358,6 +376,20 @@ def all_inputs_casadi(self):
             ]
             return self._all_inputs_casadi
 
+    @property
+    def y_last(self):
+        try:
+            return self._y_last
+        except AttributeError:
+            all_ys_last = self.all_ys[-1]
+            if isinstance(all_ys_last, pybamm.NoMemAllocVertcat):
+                self._y_last = all_ys_last.get_value()
+            elif all_ys_last.shape[1] == 1:
+                self._y_last = all_ys_last
+            else:
+                self._y_last = all_ys_last[:, -1]
+            return self._y_last
+
     @property
     def t_event(self):
         """Time at which the event happens"""
@@ -719,21 +751,30 @@ def __add__(self, other):
         # Update list of sub-solutions
         if other.all_ts[0][0] == self.all_ts[-1][-1]:
             # Skip first time step if it is repeated
-            all_ts = self.all_ts + [other.all_ts[0][1:]] + other.all_ts[1:]
-            all_ys = self.all_ys + [other.all_ys[0][:, 1:]] + other.all_ys[1:]
+            if len(other.all_ts[0]) == 1:
+                all_ts = self.all_ts + other.all_ts[1:]
+                all_ys = self.all_ys + other.all_ys[1:]
+                all_models = self.all_models + other.all_models[1:]
+                all_inputs = self.all_inputs + other.all_inputs[1:]
+            else:
+                all_ts = self.all_ts + [other.all_ts[0][1:]] + other.all_ts[1:]
+                all_ys = self.all_ys + [other.all_ys[0][:, 1:]] + other.all_ys[1:]
+                all_models = (self.all_models + other.all_models,)
+                all_inputs = (self.all_inputs + other.all_inputs,)
         else:
             all_ts = self.all_ts + other.all_ts
             all_ys = self.all_ys + other.all_ys
 
         new_sol = Solution(
             all_ts,
             all_ys,
-            self.all_models + other.all_models,
-            self.all_inputs + other.all_inputs,
+            all_models,
+            all_inputs,
             other.t_event,
             other.y_event,
             other.termination,
             bool(self.sensitivities),
+            check_solution=False,
         )
 
         new_sol.closest_event_idx = other.closest_event_idx
@@ -769,6 +810,7 @@ def copy(self):
             self.t_event,
             self.y_event,
             self.termination,
+            check_solution=False,
         )
         new_sol._all_inputs_casadi = self.all_inputs_casadi
         new_sol._sub_solutions = self.sub_solutions

pybamm/solvers/solver_utils.py

@@ -1,6 +1,7 @@
 #
 # Utility functions and classes for solvers
 #
+import numpy as np
 
 
 class NoMemAllocVertcat:
@@ -11,51 +12,18 @@ class NoMemAllocVertcat:
     def __init__(self, a, b):
         arrays = [a, b]
         self.arrays = arrays
-
-        for array in arrays:
-            if not 1 <= len(array.shape) <= 2:
-                raise ValueError("Only 1D or 2D arrays are supported")
-            self._ndim = len(array.shape)
-
         self.len_a = a.shape[0]
-        shape0 = a.shape[0] + b.shape[0]
-
-        if self._ndim == 1:
-            self._shape = (shape0,)
-            self._size = shape0
-        else:
-            if a.shape[1] != b.shape[1]:
-                raise ValueError("All arrays must have the same number of columns")
-            shape1 = a.shape[1]
-
-            self._shape = (shape0, shape1)
-            self._size = shape0 * shape1
+        self.len_b = b.shape[0]
+        self.len = self.len_a + self.len_b
+        self._shape = (self.len, 1)
 
     @property
     def shape(self):
         return self._shape
 
-    @property
-    def size(self):
-        return self._size
-
-    @property
-    def ndim(self):
-        return self._ndim
-
-    def __getitem__(self, key):
-        if self._ndim == 1 or isinstance(key, int):
-            if key < self.len_a:
-                return self.arrays[0][key]
-            else:
-                return self.arrays[1][key - self.len_a]
-
-        if key[0] == slice(None):
-            return NoMemAllocVertcat(*[arr[:, key[1]] for arr in self.arrays])
-        elif isinstance(key[0], int):
-            if key[0] < self.len_a:
-                return self.arrays[0][key[0], key[1]]
-            else:
-                return self.arrays[1][key[0] - self.len_a, key[1]]
-        else:
-            raise NotImplementedError
+    def get_value(self, out=None):
+        if out is None:
+            out = np.empty((self.len, 1))
+        out[: self.len_a] = self.arrays[0]
+        out[self.len_a :] = self.arrays[1]
+        return out