Skip to content

Commit

Permalink
clean up ProcessedVariable idaklu import (#4744)
Browse files Browse the repository at this point in the history
  • Loading branch information
MarcBerliner authored Jan 10, 2025
1 parent 3dccd52 commit 414b918
Showing 1 changed file with 7 additions and 10 deletions.
17 changes: 7 additions & 10 deletions src/pybamm/solvers/processed_variable.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
from scipy.integrate import cumulative_trapezoid
import xarray as xr
import bisect
from pybammsolvers import idaklu


class ProcessedVariable:
Expand Down Expand Up @@ -140,13 +141,13 @@ def _setup_cpp_inputs(self, t, full_range):

is_f_contiguous = _is_f_contiguous(ys)

ts = pybamm.solvers.idaklu_solver.idaklu.VectorRealtypeNdArray(ts)
ys = pybamm.solvers.idaklu_solver.idaklu.VectorRealtypeNdArray(ys)
ts = idaklu.VectorRealtypeNdArray(ts)
ys = idaklu.VectorRealtypeNdArray(ys)
if self.hermite_interpolation:
yps = pybamm.solvers.idaklu_solver.idaklu.VectorRealtypeNdArray(yps)
yps = idaklu.VectorRealtypeNdArray(yps)
else:
yps = None
inputs = pybamm.solvers.idaklu_solver.idaklu.VectorRealtypeNdArray(inputs)
inputs = idaklu.VectorRealtypeNdArray(inputs)

# Generate the serialized C++ functions only once
funcs_unique = {}
Expand All @@ -164,9 +165,7 @@ def _observe_hermite_cpp(self, t):

ts, ys, yps, funcs, inputs, _ = self._setup_cpp_inputs(t, full_range=False)
shapes = self._shape(t)
return pybamm.solvers.idaklu_solver.idaklu.observe_hermite_interp(
t, ts, ys, yps, inputs, funcs, shapes
)
return idaklu.observe_hermite_interp(t, ts, ys, yps, inputs, funcs, shapes)

def _observe_raw_cpp(self):
pybamm.logger.debug("Observing the variable raw data in C++")
Expand All @@ -176,9 +175,7 @@ def _observe_raw_cpp(self):
)
shapes = self._shape(self.t_pts)

return pybamm.solvers.idaklu_solver.idaklu.observe(
ts, ys, inputs, funcs, is_f_contiguous, shapes
)
return idaklu.observe(ts, ys, inputs, funcs, is_f_contiguous, shapes)

def _observe_raw_python(self):
raise NotImplementedError # pragma: no cover
Expand Down

0 comments on commit 414b918

Please sign in to comment.