#1996 working on nomemalloc

Author: valentinsulzer

pybamm/solvers/casadi_solver.py

@@ -309,6 +309,8 @@ def _solve_for_event(self, coarse_solution, init_event_signs):
         inputs = casadi.vertcat(*[x for x in inputs_dict.values()])
 
         def find_t_event(sol, typ):
+            sol_t = sol.all_ts[0]
+            sol_y = sol.all_ys[0]
 
             # Check most recent y to see if any events have been crossed
             if model.terminate_events_eval:
@@ -317,7 +319,7 @@ def find_t_event(sol, typ):
                     init_event_signs
                     * np.concatenate(
                         [
-                            event(sol.t[-1], y_last, inputs)
+                            event(sol_t[-1], y_last, inputs)
                             for event in model.terminate_events_eval
                         ]
                     )
@@ -353,15 +355,15 @@ def f(idx):
                         # exactly on zero, as can happen when the event switch is used
                         # (fast with events mode)
                         f_eval[idx] = (
-                            init_event_sign * event(sol.t[idx], sol.y[:, idx], inputs)
+                            init_event_sign * event(sol_t[idx], sol_y[:, idx], inputs)
                             - 1e-5
                         )
                         return f_eval[idx]
 
                 def integer_bisect():
                     a_n = 0
-                    b_n = len(sol.t) - 1
-                    for _ in range(len(sol.t)):
+                    b_n = len(sol_t) - 1
+                    for _ in range(len(sol_t)):
                         if a_n + 1 == b_n:
                             return a_n
                         m_n = (a_n + b_n) // 2
@@ -383,8 +385,8 @@ def integer_bisect():
                     # Linear interpolation between the two indices to find the root time
                     # We could do cubic interpolation here instead but it would be
                     # slower
-                    t_lower = sol.t[event_idx_lower]
-                    t_upper = sol.t[event_idx_lower + 1]
+                    t_lower = sol_t[event_idx_lower]
+                    t_upper = sol_t[event_idx_lower + 1]
                     event_lower = abs(f(event_idx_lower))
                     event_upper = abs(f(event_idx_lower + 1))
 
@@ -400,7 +402,7 @@ def integer_bisect():
                 t_event = np.nanmin(t_events)
                 # create interpolant to evaluate y in the current integration
                 # window
-                y_sol = interp1d(sol.t, sol.y, kind="linear")
+                y_sol = interp1d(sol_t, sol_y, kind="linear")
                 y_event = y_sol(t_event)
 
                 closest_event_idx = event_idx[np.nanargmin(t_events)]
@@ -433,7 +435,7 @@ def integer_bisect():
             self.create_integrator(model, inputs, t_window_event_dense)
             use_grid = True
 
-        y0 = coarse_solution.y[:, event_idx_lower]
+        y0 = coarse_solution.all_ys[0][:, event_idx_lower]
         dense_step_sol = self._run_integrator(
             model,
             y0,
@@ -458,8 +460,8 @@ def integer_bisect():
 
         # Return solution truncated at the first coarse event time
         # Also assign t_event
-        t_sol = coarse_solution.t[: event_idx_lower + 1]
-        y_sol = coarse_solution.y[:, : event_idx_lower + 1]
+        t_sol = coarse_solution.all_ts[0][: event_idx_lower + 1]
+        y_sol = coarse_solution.all_ys[0][:, : event_idx_lower + 1]
         solution = pybamm.Solution(
             t_sol,
             y_sol,
@@ -677,6 +679,7 @@ def _run_integrator(
             integrator = self.integrators[model]["no grid"]
 
         len_rhs = model.len_rhs
+        len_alg = model.len_alg
         len_t = len(t_eval)
 
         # Check y0 to see if it includes sensitivities
@@ -705,14 +708,13 @@ def _run_integrator(
                     y_sols = x_sols
                 else:
                     z_sols = casadi.horzsplit(casadi_sol["zf"])
-                    y_sols = [
-                        pybamm.NoMemAllocVertcat(x, z) for x, z in zip(x_sols, z_sols)
-                    ]
+                    y_sols = pybamm.NoMemAllocVertcat(x_sols, z_sols, len_rhs, len_alg)
             else:
                 # Repeated calls to the integrator
-                y_sols = [y0]
                 x = y0_diff
                 z = y0_alg
+                x_sols = [x]
+                z_sols = [z]
                 for i in range(len_t - 1):
                     t_min = t_eval[i]
                     t_max = t_eval[i + 1]
@@ -724,16 +726,19 @@ def _run_integrator(
                     integration_time = timer.time()
                     x = casadi_sol["xf"]
                     z = casadi_sol["zf"]
-                    if z.is_empty():
-                        y_sols.append(x)
-                    else:
-                        y_sols.append(pybamm.NoMemAllocVertcat(x, z))
+                    x_sols.append(x)
+                    if not z.is_empty():
+                        z_sols.append(z)
+                if z.is_empty():
+                    y_sols = x_sols
+                else:
+                    y_sols = pybamm.NoMemAllocVertcat(x_sols, z_sols, len_rhs, len_alg)
 
             sol = pybamm.Solution(
-                np.array_split(t_eval, len_t),
+                t_eval,
                 y_sols,
-                [model] * len_t,
-                [inputs_dict] * len_t,
+                model,
+                inputs_dict,
                 sensitivities=extract_sensitivities_in_solution,
                 check_solution=False,
             )

pybamm/solvers/processed_variable.py

@@ -68,7 +68,7 @@ def __init__(self, base_variables, base_variables_casadi, solution, warn=True):
         # Evaluate base variable at final time
         self.base_eval = self.base_variables_casadi[0](
             self.all_ts[-1][-1], self.y_last, self.all_inputs_casadi[-1]
-        ).full()
+        ).toarray()
 
         # handle 2D (in space) finite element variables differently
         if (
@@ -119,8 +119,8 @@ def initialise_0D(self):
         ):
             for inner_idx, t in enumerate(ts):
                 t = ts[inner_idx]
-                y = self.getitem(ys, idx, self.out)
-                entries[idx] = base_var_casadi(t, y, inputs).full()[0, 0]
+                y = ys[:, inner_idx]
+                entries[idx] = float(base_var_casadi(t, y, inputs))
                 idx += 1
 
         # set up interpolation
@@ -151,8 +151,8 @@ def initialise_1D(self, fixed_t=False):
         ):
             for inner_idx, t in enumerate(ts):
                 t = ts[inner_idx]
-                y = self.getitem(ys, idx, self.out)
-                entries[:, idx] = base_var_casadi(t, y, inputs).full()[:, 0]
+                y = ys[:, inner_idx]
+                entries[:, idx] = np.array(base_var_casadi(t, y, inputs).elements())
                 idx += 1
 
         # Get node and edge values
@@ -252,9 +252,9 @@ def initialise_2D(self):
         ):
             for inner_idx, t in enumerate(ts):
                 t = ts[inner_idx]
-                y = self.getitem(ys, idx, self.out)
+                y = ys[:, inner_idx]
                 entries[:, :, idx] = np.reshape(
-                    base_var_casadi(t, y, inputs).full(),
+                    np.array(base_var_casadi(t, y, inputs).elements()),
                     [first_dim_size, second_dim_size],
                     order="F",
                 )
@@ -424,9 +424,9 @@ def initialise_2D_scikit_fem(self):
         ):
             for inner_idx, t in enumerate(ts):
                 t = ts[inner_idx]
-                y = self.getitem(ys, idx, self.out)
+                y = ys[:, inner_idx]
                 entries[:, :, idx] = np.reshape(
-                    base_var_casadi(t, y, inputs).full(),
+                    base_var_casadi(t, y, inputs).toarray(),
                     [len_y, len_z],
                     order="C",
                 )
@@ -459,13 +459,6 @@ 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),
@@ -595,7 +588,7 @@ def initialise_sensitivity_explicit_forward(self):
         )
         for index, (ts, ys) in enumerate(zip(self.all_ts, self.all_ys)):
             for idx, t in enumerate(ts):
-                u = ys[:, idx]
+                u = ys[:, inner_idx]
                 next_dvar_dy_eval = dvar_dy_func(t, u, inputs_stacked)
                 next_dvar_dp_eval = dvar_dp_func(t, u, inputs_stacked)
                 if index == 0 and idx == 0:

pybamm/solvers/solution.py

@@ -97,25 +97,6 @@ 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
@@ -383,7 +364,7 @@ def y_last(self):
         except AttributeError:
             all_ys_last = self.all_ys[-1]
             if isinstance(all_ys_last, pybamm.NoMemAllocVertcat):
-                self._y_last = all_ys_last.get_value()
+                self._y_last = all_ys_last[:, -1]
             elif all_ys_last.shape[1] == 1:
                 self._y_last = all_ys_last
             else:
@@ -751,19 +732,13 @@ 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
-            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,)
+            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:]
         else:
             all_ts = self.all_ts + other.all_ts
             all_ys = self.all_ys + other.all_ys
+            all_models = self.all_models + other.all_models
+            all_inputs = self.all_inputs + other.all_inputs
 
         new_sol = Solution(
             all_ts,

pybamm/solvers/solver_utils.py

@@ -1,29 +1,43 @@
 #
 # Utility functions and classes for solvers
 #
-import numpy as np
+import casadi
 
 
 class NoMemAllocVertcat:
     """
     Acts like a vertcat, but does not allocate new memory.
     """
 
-    def __init__(self, a, b):
-        arrays = [a, b]
-        self.arrays = arrays
-        self.len_a = a.shape[0]
-        self.len_b = b.shape[0]
-        self.len = self.len_a + self.len_b
-        self._shape = (self.len, 1)
+    def __init__(self, xs, ys, len_x=None, len_y=None, items=None):
+        self.xs = xs
+        self.ys = ys
+        self.len_x = len_x or xs[0].shape[0]
+        self.len_y = len_y or ys[0].shape[0]
+        len_items = len(xs)
+        self.shape = (self.len_x + self.len_y, len_items)
 
-    @property
-    def shape(self):
-        return self._shape
+        if items is None:
+            items = [None] * len_items
+        self.items = items
 
-    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
+    def __getitem__(self, idx):
+        if idx[0] != slice(None):
+            raise NotImplementedError(
+                "Only full slices are supported in the first entry of the index"
+            )
+        idx = idx[1]
+        if isinstance(idx, slice):
+            return NoMemAllocVertcat(
+                self.xs[idx], self.ys[idx], self.len_x, self.len_y, self.items[idx]
+            )
+        else:
+            item = self.items[idx]
+            if item is not None:
+                return item
+            else:
+                out = casadi.DM.zeros((self.shape[0], 1))
+                out[: self.len_x] = self.xs[idx]
+                out[self.len_x :] = self.ys[idx]
+                self.items[idx] = out
+                return out