def test_numpy_concatenation_vector_scalar(self): # with entries y = np.linspace(0, 1, 10)[:, np.newaxis] a = pybamm.Vector(y) b = pybamm.Scalar(16) c = pybamm.Scalar(3) conc = pybamm.NumpyConcatenation(a, b, c) np.testing.assert_array_equal( conc.evaluate(y=y), np.concatenate([y, np.array([[16]]), np.array([[3]])])) # with y_slice a = pybamm.StateVector(slice(0, 10)) conc = pybamm.NumpyConcatenation(a, b, c) np.testing.assert_array_equal( conc.evaluate(y=y), np.concatenate([y, np.array([[16]]), np.array([[3]])])) # with time b = pybamm.t conc = pybamm.NumpyConcatenation(a, b, c) np.testing.assert_array_equal( conc.evaluate(16, y), np.concatenate([y, np.array([[16]]), np.array([[3]])]))
def test_jac_of_numpy_concatenation(self): u = pybamm.StateVector(slice(0, 2)) y0 = np.ones(2) # Multiple children func = pybamm.NumpyConcatenation(u, u) jacobian = np.array([[1, 0], [0, 1], [1, 0], [0, 1]]) dfunc_dy = func.jac(u).evaluate(y=y0) np.testing.assert_array_equal(jacobian, dfunc_dy.toarray()) # One child self.assertEqual(u.jac(u).id, pybamm.NumpyConcatenation(u).jac(u).id)
def test_numpy_concatenation_vectors(self): # with entries y = np.linspace(0, 1, 15)[:, np.newaxis] a = pybamm.Vector(y[:5]) b = pybamm.Vector(y[5:9]) c = pybamm.Vector(y[9:]) conc = pybamm.NumpyConcatenation(a, b, c) np.testing.assert_array_equal(conc.evaluate(None, y), y) # with y_slice a = pybamm.StateVector(slice(0, 10)) b = pybamm.StateVector(slice(10, 15)) c = pybamm.StateVector(slice(15, 23)) conc = pybamm.NumpyConcatenation(a, b, c) y = np.linspace(0, 1, 23)[:, np.newaxis] np.testing.assert_array_equal(conc.evaluate(None, y), y)
def test_concatenations(self): y = np.linspace(0, 1, 10)[:, np.newaxis] a = pybamm.Vector(y) b = pybamm.Scalar(16) c = pybamm.Scalar(3) conc = pybamm.NumpyConcatenation(a, b, c) self.assert_casadi_equal(conc.to_casadi(), casadi.MX(conc.evaluate()), evalf=True) # Domain concatenation mesh = get_mesh_for_testing() a_dom = ["negative electrode"] b_dom = ["separator"] a = 2 * pybamm.Vector(np.ones_like(mesh[a_dom[0]].nodes), domain=a_dom) b = pybamm.Vector(np.ones_like(mesh[b_dom[0]].nodes), domain=b_dom) conc = pybamm.DomainConcatenation([b, a], mesh) self.assert_casadi_equal(conc.to_casadi(), casadi.MX(conc.evaluate()), evalf=True) # 2d disc = get_1p1d_discretisation_for_testing() a = pybamm.Variable("a", domain=a_dom) b = pybamm.Variable("b", domain=b_dom) conc = pybamm.Concatenation(a, b) disc.set_variable_slices([conc]) expr = disc.process_symbol(conc) y = casadi.SX.sym("y", expr.size) x = expr.to_casadi(None, y) f = casadi.Function("f", [x], [x]) y_eval = np.linspace(0, 1, expr.size) self.assert_casadi_equal(f(y_eval), casadi.SX(expr.evaluate(y=y_eval)))
def test_numpy_concatenation(self): a = pybamm.Variable("a") b = pybamm.Variable("b") c = pybamm.Variable("c") self.assertEqual( pybamm.numpy_concatenation(pybamm.numpy_concatenation(a, b), c).id, pybamm.NumpyConcatenation(a, b, c).id, )
def test_symbol_new_copy(self): a = pybamm.Parameter("a") b = pybamm.Parameter("b") v_n = pybamm.Variable("v", "negative electrode") x_n = pybamm.standard_spatial_vars.x_n v_s = pybamm.Variable("v", "separator") vec = pybamm.Vector([1, 2, 3, 4, 5]) mat = pybamm.Matrix([[1, 2], [3, 4]]) mesh = get_mesh_for_testing() for symbol in [ a + b, a - b, a * b, a / b, a**b, -a, abs(a), pybamm.Function(np.sin, a), pybamm.FunctionParameter("function", {"a": a}), pybamm.grad(v_n), pybamm.div(pybamm.grad(v_n)), pybamm.upwind(v_n), pybamm.IndefiniteIntegral(v_n, x_n), pybamm.BackwardIndefiniteIntegral(v_n, x_n), pybamm.BoundaryValue(v_n, "right"), pybamm.BoundaryGradient(v_n, "right"), pybamm.PrimaryBroadcast(a, "domain"), pybamm.SecondaryBroadcast(v_n, "current collector"), pybamm.FullBroadcast(a, "domain", {"secondary": "other domain"}), pybamm.concatenation(v_n, v_s), pybamm.NumpyConcatenation(a, b, v_s), pybamm.DomainConcatenation([v_n, v_s], mesh), pybamm.Parameter("param"), pybamm.InputParameter("param"), pybamm.StateVector(slice(0, 56)), pybamm.Matrix(np.ones((50, 40))), pybamm.SpatialVariable("x", ["negative electrode"]), pybamm.t, pybamm.Index(vec, 1), pybamm.NotConstant(a), pybamm.ExternalVariable( "external variable", 20, domain="test", auxiliary_domains={"secondary": "test2"}, ), pybamm.minimum(a, b), pybamm.maximum(a, b), pybamm.SparseStack(mat, mat), ]: self.assertEqual(symbol.id, symbol.new_copy().id)
def test_symbol_new_copy(self): a = pybamm.Scalar(0) b = pybamm.Scalar(1) v_n = pybamm.Variable("v", "negative electrode") x_n = pybamm.standard_spatial_vars.x_n v_s = pybamm.Variable("v", "separator") vec = pybamm.Vector(np.array([1, 2, 3, 4, 5])) mesh = get_mesh_for_testing() for symbol in [ a + b, a - b, a * b, a / b, a**b, -a, abs(a), pybamm.Function(np.sin, a), pybamm.FunctionParameter("function", {"a": a}), pybamm.grad(v_n), pybamm.div(pybamm.grad(v_n)), pybamm.Integral(a, pybamm.t), pybamm.IndefiniteIntegral(v_n, x_n), pybamm.BackwardIndefiniteIntegral(v_n, x_n), pybamm.BoundaryValue(v_n, "right"), pybamm.BoundaryGradient(v_n, "right"), pybamm.PrimaryBroadcast(a, "domain"), pybamm.SecondaryBroadcast(v_n, "current collector"), pybamm.FullBroadcast(a, "domain", {"secondary": "other domain"}), pybamm.Concatenation(v_n, v_s), pybamm.NumpyConcatenation(a, b, v_s), pybamm.DomainConcatenation([v_n, v_s], mesh), pybamm.Parameter("param"), pybamm.InputParameter("param"), pybamm.StateVector(slice(0, 56)), pybamm.Matrix(np.ones((50, 40))), pybamm.SpatialVariable("x", ["negative electrode"]), pybamm.t, pybamm.Index(vec, 1), ]: self.assertEqual(symbol.id, symbol.new_copy().id)
def test_numpy_concatenation_simplify(self): a = pybamm.Variable("a") b = pybamm.Variable("b") c = pybamm.Variable("c") # simplifying flattens the concatenations into a single concatenation self.assertEqual( pybamm.NumpyConcatenation(pybamm.NumpyConcatenation(a, b), c).simplify().id, pybamm.NumpyConcatenation(a, b, c).id, ) self.assertEqual( pybamm.NumpyConcatenation(a, pybamm.NumpyConcatenation( b, c)).simplify().id, pybamm.NumpyConcatenation(a, b, c).id, )
def test_model_solver_manually_update_initial_conditions(self): # Create model model = pybamm.BaseModel() var1 = pybamm.Variable("var1") model.rhs = {var1: -var1} model.initial_conditions = {var1: 1} # Solve solver = pybamm.ScipySolver(rtol=1e-8, atol=1e-8) t_eval = np.linspace(0, 5, 100) solution = solver.solve(model, t_eval) np.testing.assert_array_almost_equal(solution.y[0], 1 * np.exp(-solution.t), decimal=5) # Change initial conditions and solve again model.concatenated_initial_conditions = pybamm.NumpyConcatenation( pybamm.Vector([[2]])) solution = solver.solve(model, t_eval) np.testing.assert_array_almost_equal(solution.y[0], 2 * np.exp(-solution.t), decimal=5)
def test_evaluator_python(self): a = pybamm.StateVector(slice(0, 1)) b = pybamm.StateVector(slice(1, 2)) y_tests = [np.array([[2], [3]]), np.array([[1], [3]])] t_tests = [1, 2] # test a * b expr = a * b evaluator = pybamm.EvaluatorPython(expr) result = evaluator.evaluate(t=None, y=np.array([[2], [3]])) self.assertEqual(result, 6) result = evaluator.evaluate(t=None, y=np.array([[1], [3]])) self.assertEqual(result, 3) # test function(a*b) expr = pybamm.Function(test_function, a * b) evaluator = pybamm.EvaluatorPython(expr) result = evaluator.evaluate(t=None, y=np.array([[2], [3]])) self.assertEqual(result, 12) # test a constant expression expr = pybamm.Scalar(2) * pybamm.Scalar(3) evaluator = pybamm.EvaluatorPython(expr) result = evaluator.evaluate() self.assertEqual(result, 6) # test a larger expression expr = a * b + b + a**2 / b + 2 * a + b / 2 + 4 evaluator = pybamm.EvaluatorPython(expr) for y in y_tests: result = evaluator.evaluate(t=None, y=y) self.assertEqual(result, expr.evaluate(t=None, y=y)) # test something with time expr = a * pybamm.t evaluator = pybamm.EvaluatorPython(expr) for t, y in zip(t_tests, y_tests): result = evaluator.evaluate(t=t, y=y) self.assertEqual(result, expr.evaluate(t=t, y=y)) # test something with a matrix multiplication A = pybamm.Matrix(np.array([[1, 2], [3, 4]])) expr = A @ pybamm.StateVector(slice(0, 2)) evaluator = pybamm.EvaluatorPython(expr) for t, y in zip(t_tests, y_tests): result = evaluator.evaluate(t=t, y=y) np.testing.assert_allclose(result, expr.evaluate(t=t, y=y)) # test something with a heaviside a = pybamm.Vector(np.array([1, 2])) expr = a <= pybamm.StateVector(slice(0, 2)) evaluator = pybamm.EvaluatorPython(expr) for t, y in zip(t_tests, y_tests): result = evaluator.evaluate(t=t, y=y) np.testing.assert_allclose(result, expr.evaluate(t=t, y=y)) expr = a > pybamm.StateVector(slice(0, 2)) evaluator = pybamm.EvaluatorPython(expr) for t, y in zip(t_tests, y_tests): result = evaluator.evaluate(t=t, y=y) np.testing.assert_allclose(result, expr.evaluate(t=t, y=y)) # test something with a minimum or maximum a = pybamm.Vector(np.array([1, 2])) expr = pybamm.minimum(a, pybamm.StateVector(slice(0, 2))) evaluator = pybamm.EvaluatorPython(expr) for t, y in zip(t_tests, y_tests): result = evaluator.evaluate(t=t, y=y) np.testing.assert_allclose(result, expr.evaluate(t=t, y=y)) expr = pybamm.maximum(a, pybamm.StateVector(slice(0, 2))) evaluator = pybamm.EvaluatorPython(expr) for t, y in zip(t_tests, y_tests): result = evaluator.evaluate(t=t, y=y) np.testing.assert_allclose(result, expr.evaluate(t=t, y=y)) # test something with an index expr = pybamm.Index(A @ pybamm.StateVector(slice(0, 2)), 0) evaluator = pybamm.EvaluatorPython(expr) for t, y in zip(t_tests, y_tests): result = evaluator.evaluate(t=t, y=y) self.assertEqual(result, expr.evaluate(t=t, y=y)) # test something with a sparse matrix multiplication A = pybamm.Matrix(np.array([[1, 2], [3, 4]])) B = pybamm.Matrix(scipy.sparse.csr_matrix(np.array([[1, 0], [0, 4]]))) C = pybamm.Matrix(scipy.sparse.coo_matrix(np.array([[1, 0], [0, 4]]))) expr = A @ B @ C @ pybamm.StateVector(slice(0, 2)) evaluator = pybamm.EvaluatorPython(expr) for t, y in zip(t_tests, y_tests): result = evaluator.evaluate(t=t, y=y) np.testing.assert_allclose(result, expr.evaluate(t=t, y=y)) # test numpy concatenation a = pybamm.Vector(np.array([[1], [2]])) b = pybamm.Vector(np.array([[3]])) expr = pybamm.NumpyConcatenation(a, b) evaluator = pybamm.EvaluatorPython(expr) for t, y in zip(t_tests, y_tests): result = evaluator.evaluate(t=t, y=y) np.testing.assert_allclose(result, expr.evaluate(t=t, y=y)) # test sparse stack A = pybamm.Matrix(scipy.sparse.csr_matrix(np.array([[1, 0], [0, 4]]))) B = pybamm.Matrix(scipy.sparse.csr_matrix(np.array([[2, 0], [5, 0]]))) expr = pybamm.SparseStack(A, B) evaluator = pybamm.EvaluatorPython(expr) for t, y in zip(t_tests, y_tests): result = evaluator.evaluate(t=t, y=y).toarray() np.testing.assert_allclose(result, expr.evaluate(t=t, y=y).toarray()) # test Inner v = pybamm.Vector(np.ones(5), domain="test") w = pybamm.Vector(2 * np.ones(5), domain="test") expr = pybamm.Inner(v, w) evaluator = pybamm.EvaluatorPython(expr) result = evaluator.evaluate() np.testing.assert_allclose(result, expr.evaluate())
def test_find_symbols(self): a = pybamm.StateVector(slice(0, 1)) b = pybamm.StateVector(slice(1, 2)) # test a + b constant_symbols = OrderedDict() variable_symbols = OrderedDict() expr = a + b pybamm.find_symbols(expr, constant_symbols, variable_symbols) self.assertEqual(len(constant_symbols), 0) # test keys of known_symbols self.assertEqual(list(variable_symbols.keys())[0], a.id) self.assertEqual(list(variable_symbols.keys())[1], b.id) self.assertEqual(list(variable_symbols.keys())[2], expr.id) # test values of variable_symbols self.assertEqual(list(variable_symbols.values())[0], "y[:1][[True]]") self.assertEqual( list(variable_symbols.values())[1], "y[:2][[False, True]]") var_a = pybamm.id_to_python_variable(a.id) var_b = pybamm.id_to_python_variable(b.id) self.assertEqual( list(variable_symbols.values())[2], "{} + {}".format(var_a, var_b)) # test identical subtree constant_symbols = OrderedDict() variable_symbols = OrderedDict() expr = a + b + b pybamm.find_symbols(expr, constant_symbols, variable_symbols) self.assertEqual(len(constant_symbols), 0) # test keys of variable_symbols self.assertEqual(list(variable_symbols.keys())[0], a.id) self.assertEqual(list(variable_symbols.keys())[1], b.id) self.assertEqual(list(variable_symbols.keys())[2], expr.children[0].id) self.assertEqual(list(variable_symbols.keys())[3], expr.id) # test values of variable_symbols self.assertEqual(list(variable_symbols.values())[0], "y[:1][[True]]") self.assertEqual( list(variable_symbols.values())[1], "y[:2][[False, True]]") self.assertEqual( list(variable_symbols.values())[2], "{} + {}".format(var_a, var_b)) var_child = pybamm.id_to_python_variable(expr.children[0].id) self.assertEqual( list(variable_symbols.values())[3], "{} + {}".format(var_child, var_b)) # test unary op constant_symbols = OrderedDict() variable_symbols = OrderedDict() expr = a + (-b) pybamm.find_symbols(expr, constant_symbols, variable_symbols) self.assertEqual(len(constant_symbols), 0) # test keys of variable_symbols self.assertEqual(list(variable_symbols.keys())[0], a.id) self.assertEqual(list(variable_symbols.keys())[1], b.id) self.assertEqual(list(variable_symbols.keys())[2], expr.children[1].id) self.assertEqual(list(variable_symbols.keys())[3], expr.id) # test values of variable_symbols self.assertEqual(list(variable_symbols.values())[0], "y[:1][[True]]") self.assertEqual( list(variable_symbols.values())[1], "y[:2][[False, True]]") self.assertEqual( list(variable_symbols.values())[2], "-{}".format(var_b)) var_child = pybamm.id_to_python_variable(expr.children[1].id) self.assertEqual( list(variable_symbols.values())[3], "{} + {}".format(var_a, var_child)) # test function constant_symbols = OrderedDict() variable_symbols = OrderedDict() expr = pybamm.Function(test_function, a) pybamm.find_symbols(expr, constant_symbols, variable_symbols) self.assertEqual(list(constant_symbols.keys())[0], expr.id) self.assertEqual(list(constant_symbols.values())[0], test_function) self.assertEqual(list(variable_symbols.keys())[0], a.id) self.assertEqual(list(variable_symbols.keys())[1], expr.id) self.assertEqual(list(variable_symbols.values())[0], "y[:1][[True]]") var_funct = pybamm.id_to_python_variable(expr.id, True) self.assertEqual( list(variable_symbols.values())[1], "{}({})".format(var_funct, var_a)) # test matrix constant_symbols = OrderedDict() variable_symbols = OrderedDict() A = pybamm.Matrix(np.array([[1, 2], [3, 4]])) pybamm.find_symbols(A, constant_symbols, variable_symbols) self.assertEqual(len(variable_symbols), 0) self.assertEqual(list(constant_symbols.keys())[0], A.id) np.testing.assert_allclose( list(constant_symbols.values())[0], np.array([[1, 2], [3, 4]])) # test sparse matrix constant_symbols = OrderedDict() variable_symbols = OrderedDict() A = pybamm.Matrix(scipy.sparse.csr_matrix(np.array([[0, 2], [0, 4]]))) pybamm.find_symbols(A, constant_symbols, variable_symbols) self.assertEqual(len(variable_symbols), 0) self.assertEqual(list(constant_symbols.keys())[0], A.id) np.testing.assert_allclose( list(constant_symbols.values())[0].toarray(), A.entries.toarray()) # test numpy concatentate constant_symbols = OrderedDict() variable_symbols = OrderedDict() expr = pybamm.NumpyConcatenation(a, b) pybamm.find_symbols(expr, constant_symbols, variable_symbols) self.assertEqual(len(constant_symbols), 0) self.assertEqual(list(variable_symbols.keys())[0], a.id) self.assertEqual(list(variable_symbols.keys())[1], b.id) self.assertEqual(list(variable_symbols.keys())[2], expr.id) self.assertEqual( list(variable_symbols.values())[2], "np.concatenate(({},{}))".format(var_a, var_b), ) # test domain concatentate constant_symbols = OrderedDict() variable_symbols = OrderedDict() expr = pybamm.NumpyConcatenation(a, b) pybamm.find_symbols(expr, constant_symbols, variable_symbols) self.assertEqual(len(constant_symbols), 0) self.assertEqual(list(variable_symbols.keys())[0], a.id) self.assertEqual(list(variable_symbols.keys())[1], b.id) self.assertEqual(list(variable_symbols.keys())[2], expr.id) self.assertEqual( list(variable_symbols.values())[2], "np.concatenate(({},{}))".format(var_a, var_b), ) # test that Concatentation throws expr = pybamm.Concatenation(a, b) with self.assertRaises(NotImplementedError): pybamm.find_symbols(expr, constant_symbols, variable_symbols) # test that these nodes throw for expr in (pybamm.Variable("a"), pybamm.Parameter("a")): with self.assertRaises(NotImplementedError): pybamm.find_symbols(expr, constant_symbols, variable_symbols)
def set_initial_conditions_from(self, solution, inplace=True): """ Update initial conditions with the final states from a Solution object or from a dictionary. This assumes that, for each variable in self.initial_conditions, there is a corresponding variable in the solution with the same name and size. Parameters ---------- solution : :class:`pybamm.Solution`, or dict The solution to use to initialize the model inplace : bool Whether to modify the model inplace or create a new model """ if inplace is True: model = self else: model = self.new_copy() if isinstance(solution, pybamm.Solution): solution = solution.last_state for var, equation in model.initial_conditions.items(): if isinstance(var, pybamm.Variable): try: final_state = solution[var.name] except KeyError as e: raise pybamm.ModelError( "To update a model from a solution, each variable in " "model.initial_conditions must appear in the solution with " "the same key as the variable name. In the solution provided, " f"{e.args[0]}") if isinstance(solution, pybamm.Solution): final_state = final_state.data if final_state.ndim == 1: final_state_eval = final_state[-1:] elif final_state.ndim == 2: final_state_eval = final_state[:, -1] elif final_state.ndim == 3: final_state_eval = final_state[:, :, -1].flatten(order="F") else: raise NotImplementedError("Variable must be 0D, 1D, or 2D") model.initial_conditions[var] = pybamm.Vector(final_state_eval) elif isinstance(var, pybamm.Concatenation): children = [] for child in var.orphans: try: final_state = solution[child.name] except KeyError as e: raise pybamm.ModelError( "To update a model from a solution, each variable in " "model.initial_conditions must appear in the solution with " "the same key as the variable name. In the solution " f"provided, {e.args[0]}") if isinstance(solution, pybamm.Solution): final_state = final_state.data if final_state.ndim == 2: final_state_eval = final_state[:, -1] else: raise NotImplementedError( "Variable in concatenation must be 1D") children.append(final_state_eval) model.initial_conditions[var] = pybamm.Vector( np.concatenate(children)) else: raise NotImplementedError( "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 concatenate(self, *symbols, sparse=False): if sparse: return pybamm.SparseStack(*symbols) else: return pybamm.NumpyConcatenation(*symbols)
def set_up(self, model, inputs=None): """Unpack model, perform checks, simplify and calculate jacobian. Parameters ---------- model : :class:`pybamm.BaseModel` The model whose solution to calculate. Must have attributes rhs and initial_conditions inputs : dict, optional Any input parameters to pass to the model when solving """ inputs = inputs or {} y0 = model.concatenated_initial_conditions.evaluate(0, None, inputs) # Set model timescale model.timescale_eval = model.timescale.evaluate(u=inputs) # Check model.algebraic for ode solvers if self.ode_solver is True and len(model.algebraic) > 0: raise pybamm.SolverError( "Cannot use ODE solver '{}' to solve DAE model".format( self.name)) if self.ode_solver is True: self.root_method = None if (isinstance(self, pybamm.CasadiSolver) or self.root_method == "casadi") and model.convert_to_format != "casadi": pybamm.logger.warning( f"Converting {model.name} to CasADi for solving with CasADi solver" ) model.convert_to_format = "casadi" if model.convert_to_format != "casadi": simp = pybamm.Simplification() # Create Jacobian from concatenated rhs and algebraic y = pybamm.StateVector(slice(0, np.size(y0))) # set up Jacobian object, for re-use of dict jacobian = pybamm.Jacobian() else: # Convert model attributes to casadi t_casadi = casadi.MX.sym("t") y_diff = casadi.MX.sym( "y_diff", len(model.concatenated_rhs.evaluate(0, y0, inputs))) y_alg = casadi.MX.sym( "y_alg", len(model.concatenated_algebraic.evaluate(0, y0, inputs))) y_casadi = casadi.vertcat(y_diff, y_alg) u_casadi = {} for name, value in inputs.items(): if isinstance(value, numbers.Number): u_casadi[name] = casadi.MX.sym(name) else: u_casadi[name] = casadi.MX.sym(name, value.shape[0]) u_casadi_stacked = casadi.vertcat(*[u for u in u_casadi.values()]) def process(func, name, use_jacobian=None): def report(string): # don't log event conversion if "event" not in string: pybamm.logger.info(string) if use_jacobian is None: use_jacobian = model.use_jacobian if model.convert_to_format != "casadi": # Process with pybamm functions if model.use_simplify: report(f"Simplifying {name}") func = simp.simplify(func) if use_jacobian: report(f"Calculating jacobian for {name}") jac = jacobian.jac(func, y) if model.use_simplify: report(f"Simplifying jacobian for {name}") jac = simp.simplify(jac) if model.convert_to_format == "python": report(f"Converting jacobian for {name} to python") jac = pybamm.EvaluatorPython(jac) jac = jac.evaluate else: jac = None if model.convert_to_format == "python": report(f"Converting {name} to python") func = pybamm.EvaluatorPython(func) func = func.evaluate else: # Process with CasADi report(f"Converting {name} to CasADi") func = func.to_casadi(t_casadi, y_casadi, u_casadi) if use_jacobian: report(f"Calculating jacobian for {name} using CasADi") jac_casadi = casadi.jacobian(func, y_casadi) jac = casadi.Function( name, [t_casadi, y_casadi, u_casadi_stacked], [jac_casadi]) else: jac = None func = casadi.Function(name, [t_casadi, y_casadi, u_casadi_stacked], [func]) if name == "residuals": func_call = Residuals(func, name, model) else: func_call = SolverCallable(func, name, model) func_call.set_inputs(inputs) if jac is not None: jac_call = SolverCallable(jac, name + "_jac", model) jac_call.set_inputs(inputs) else: jac_call = None return func, func_call, jac_call # Check for heaviside functions in rhs and algebraic and add discontinuity # events if these exist. # Note: only checks for the case of t < X, t <= X, X < t, or X <= t, but also # accounts for the fact that t might be dimensional # Only do this for DAE models as ODE models can deal with discontinuities fine if len(model.algebraic) > 0: for symbol in itertools.chain( model.concatenated_rhs.pre_order(), model.concatenated_algebraic.pre_order(), ): if isinstance(symbol, pybamm.Heaviside): # Dimensionless if symbol.right.id == pybamm.t.id: expr = symbol.left elif symbol.left.id == pybamm.t.id: expr = symbol.right # Dimensional elif symbol.right.id == (pybamm.t * model.timescale).id: expr = symbol.left.new_copy( ) / symbol.right.right.new_copy() elif symbol.left.id == (pybamm.t * model.timescale).id: expr = symbol.right.new_copy( ) / symbol.left.right.new_copy() model.events.append( pybamm.Event(str(symbol), expr.new_copy(), pybamm.EventType.DISCONTINUITY)) # Process rhs, algebraic and event expressions rhs, rhs_eval, jac_rhs = process(model.concatenated_rhs, "RHS") algebraic, algebraic_eval, jac_algebraic = process( model.concatenated_algebraic, "algebraic") terminate_events_eval = [ process(event.expression, "event", use_jacobian=False)[1] for event in model.events if event.event_type == pybamm.EventType.TERMINATION ] # discontinuity events are evaluated before the solver is called, so don't need # to process them discontinuity_events_eval = [ event for event in model.events if event.event_type == pybamm.EventType.DISCONTINUITY ] # Add the solver attributes model.rhs_eval = rhs_eval model.algebraic_eval = algebraic_eval model.jac_algebraic_eval = jac_algebraic model.terminate_events_eval = terminate_events_eval model.discontinuity_events_eval = discontinuity_events_eval # Save CasADi functions for the CasADi solver # Note: when we pass to casadi the ode part of the problem must be in explicit # form so we pre-multiply by the inverse of the mass matrix if self.root_method == "casadi" or isinstance(self, pybamm.CasadiSolver): mass_matrix_inv = casadi.MX(model.mass_matrix_inv.entries) explicit_rhs = mass_matrix_inv @ rhs(t_casadi, y_casadi, u_casadi_stacked) model.casadi_rhs = casadi.Function( "rhs", [t_casadi, y_casadi, u_casadi_stacked], [explicit_rhs]) model.casadi_algebraic = algebraic # Calculate consistent initial conditions for the algebraic equations if len(model.algebraic) > 0: all_states = pybamm.NumpyConcatenation( model.concatenated_rhs, model.concatenated_algebraic) # Process again, uses caching so should be quick residuals, residuals_eval, jacobian_eval = process( all_states, "residuals") model.residuals_eval = residuals_eval model.jacobian_eval = jacobian_eval y0_guess = y0.flatten() model.y0 = self.calculate_consistent_state(model, 0, y0_guess, inputs) else: # can use DAE solver to solve ODE model model.residuals_eval = Residuals(rhs, "residuals", model) model.jacobian_eval = jac_rhs model.y0 = y0.flatten() pybamm.logger.info("Finish solver set-up")
def set_up(self, model, inputs=None, t_eval=None): """Unpack model, perform checks, simplify and calculate jacobian. Parameters ---------- model : :class:`pybamm.BaseModel` The model whose solution to calculate. Must have attributes rhs and initial_conditions inputs : dict, optional Any input parameters to pass to the model when solving t_eval : numeric type, optional The times (in seconds) at which to compute the solution """ # Check model.algebraic for ode solvers if self.ode_solver is True and len(model.algebraic) > 0: raise pybamm.SolverError( "Cannot use ODE solver '{}' to solve DAE model".format( self.name)) # Check model.rhs for algebraic solvers if self.algebraic_solver is True and len(model.rhs) > 0: raise pybamm.SolverError( """Cannot use algebraic solver to solve model with time derivatives""" ) # casadi solver won't allow solving algebraic model so we have to raise an # error here if isinstance(self, pybamm.CasadiSolver) and len(model.rhs) == 0: raise pybamm.SolverError( "Cannot use CasadiSolver to solve algebraic model, " "use CasadiAlgebraicSolver instead") # Discretise model if it isn't already discretised # This only works with purely 0D models, as otherwise the mesh and spatial # method should be specified by the user if model.is_discretised is False: try: disc = pybamm.Discretisation() disc.process_model(model) except pybamm.DiscretisationError as e: raise pybamm.DiscretisationError( "Cannot automatically discretise model, " "model should be discretised before solving ({})".format( e)) inputs = inputs or {} # Set model timescale model.timescale_eval = model.timescale.evaluate(inputs=inputs) # Set model lengthscales model.length_scales_eval = { domain: scale.evaluate(inputs=inputs) for domain, scale in model.length_scales.items() } if (isinstance(self, (pybamm.CasadiSolver, pybamm.CasadiAlgebraicSolver)) ) and model.convert_to_format != "casadi": pybamm.logger.warning( "Converting {} to CasADi for solving with CasADi solver". format(model.name)) model.convert_to_format = "casadi" if (isinstance(self.root_method, pybamm.CasadiAlgebraicSolver) and model.convert_to_format != "casadi"): pybamm.logger.warning( "Converting {} to CasADi for calculating ICs with CasADi". format(model.name)) model.convert_to_format = "casadi" if model.convert_to_format != "casadi": simp = pybamm.Simplification() # Create Jacobian from concatenated rhs and algebraic y = pybamm.StateVector( slice(0, model.concatenated_initial_conditions.size)) # set up Jacobian object, for re-use of dict jacobian = pybamm.Jacobian() else: # Convert model attributes to casadi t_casadi = casadi.MX.sym("t") y_diff = casadi.MX.sym("y_diff", model.concatenated_rhs.size) y_alg = casadi.MX.sym("y_alg", model.concatenated_algebraic.size) y_casadi = casadi.vertcat(y_diff, y_alg) p_casadi = {} for name, value in inputs.items(): if isinstance(value, numbers.Number): p_casadi[name] = casadi.MX.sym(name) else: p_casadi[name] = casadi.MX.sym(name, value.shape[0]) p_casadi_stacked = casadi.vertcat(*[p for p in p_casadi.values()]) def process(func, name, use_jacobian=None): def report(string): # don't log event conversion if "event" not in string: pybamm.logger.info(string) if use_jacobian is None: use_jacobian = model.use_jacobian if model.convert_to_format != "casadi": # Process with pybamm functions if model.use_simplify: report(f"Simplifying {name}") func = simp.simplify(func) if model.convert_to_format == "jax": report(f"Converting {name} to jax") jax_func = pybamm.EvaluatorJax(func) if use_jacobian: report(f"Calculating jacobian for {name}") jac = jacobian.jac(func, y) if model.use_simplify: report(f"Simplifying jacobian for {name}") jac = simp.simplify(jac) if model.convert_to_format == "python": report(f"Converting jacobian for {name} to python") jac = pybamm.EvaluatorPython(jac) elif model.convert_to_format == "jax": report(f"Converting jacobian for {name} to jax") jac = jax_func.get_jacobian() jac = jac.evaluate else: jac = None if model.convert_to_format == "python": report(f"Converting {name} to python") func = pybamm.EvaluatorPython(func) if model.convert_to_format == "jax": report(f"Converting {name} to jax") func = jax_func func = func.evaluate else: # Process with CasADi report(f"Converting {name} to CasADi") func = func.to_casadi(t_casadi, y_casadi, inputs=p_casadi) if use_jacobian: report(f"Calculating jacobian for {name} using CasADi") jac_casadi = casadi.jacobian(func, y_casadi) jac = casadi.Function( name, [t_casadi, y_casadi, p_casadi_stacked], [jac_casadi]) else: jac = None func = casadi.Function(name, [t_casadi, y_casadi, p_casadi_stacked], [func]) if name == "residuals": func_call = Residuals(func, name, model) else: func_call = SolverCallable(func, name, model) if jac is not None: jac_call = SolverCallable(jac, name + "_jac", model) else: jac_call = None return func, func_call, jac_call # Check for heaviside and modulo functions in rhs and algebraic and add # discontinuity events if these exist. # Note: only checks for the case of t < X, t <= X, X < t, or X <= t, but also # accounts for the fact that t might be dimensional # Only do this for DAE models as ODE models can deal with discontinuities fine if len(model.algebraic) > 0: for symbol in itertools.chain( model.concatenated_rhs.pre_order(), model.concatenated_algebraic.pre_order(), ): if isinstance(symbol, pybamm.Heaviside): found_t = False # Dimensionless if symbol.right.id == pybamm.t.id: expr = symbol.left found_t = True elif symbol.left.id == pybamm.t.id: expr = symbol.right found_t = True # Dimensional elif symbol.right.id == (pybamm.t * model.timescale).id: expr = symbol.left.new_copy( ) / symbol.right.right.new_copy() found_t = True elif symbol.left.id == (pybamm.t * model.timescale).id: expr = symbol.right.new_copy( ) / symbol.left.right.new_copy() found_t = True # Update the events if the heaviside function depended on t if found_t: model.events.append( pybamm.Event( str(symbol), expr.new_copy(), pybamm.EventType.DISCONTINUITY, )) elif isinstance(symbol, pybamm.Modulo): found_t = False # Dimensionless if symbol.left.id == pybamm.t.id: expr = symbol.right found_t = True # Dimensional elif symbol.left.id == (pybamm.t * model.timescale).id: expr = symbol.right.new_copy( ) / symbol.left.right.new_copy() found_t = True # Update the events if the modulo function depended on t if found_t: if t_eval is None: N_events = 200 else: N_events = t_eval[-1] // expr.value for i in np.arange(N_events): model.events.append( pybamm.Event( str(symbol), expr.new_copy() * pybamm.Scalar(i + 1), pybamm.EventType.DISCONTINUITY, )) # Process initial conditions initial_conditions = process( model.concatenated_initial_conditions, "initial_conditions", use_jacobian=False, )[0] init_eval = InitialConditions(initial_conditions, model) # Process rhs, algebraic and event expressions rhs, rhs_eval, jac_rhs = process(model.concatenated_rhs, "RHS") algebraic, algebraic_eval, jac_algebraic = process( model.concatenated_algebraic, "algebraic") terminate_events_eval = [ process(event.expression, "event", use_jacobian=False)[1] for event in model.events if event.event_type == pybamm.EventType.TERMINATION ] # discontinuity events are evaluated before the solver is called, so don't need # to process them discontinuity_events_eval = [ event for event in model.events if event.event_type == pybamm.EventType.DISCONTINUITY ] # Add the solver attributes model.init_eval = init_eval model.rhs_eval = rhs_eval model.algebraic_eval = algebraic_eval model.jac_algebraic_eval = jac_algebraic model.terminate_events_eval = terminate_events_eval model.discontinuity_events_eval = discontinuity_events_eval # Calculate initial conditions model.y0 = init_eval(inputs) # Save CasADi functions for the CasADi solver # Note: when we pass to casadi the ode part of the problem must be in explicit # form so we pre-multiply by the inverse of the mass matrix if isinstance( self.root_method, pybamm.CasadiAlgebraicSolver) or isinstance( self, (pybamm.CasadiSolver, pybamm.CasadiAlgebraicSolver)): # can use DAE solver to solve model with algebraic equations only if len(model.rhs) > 0: mass_matrix_inv = casadi.MX(model.mass_matrix_inv.entries) explicit_rhs = mass_matrix_inv @ rhs(t_casadi, y_casadi, p_casadi_stacked) model.casadi_rhs = casadi.Function( "rhs", [t_casadi, y_casadi, p_casadi_stacked], [explicit_rhs]) model.casadi_algebraic = algebraic if len(model.rhs) == 0: # No rhs equations: residuals is algebraic only model.residuals_eval = Residuals(algebraic, "residuals", model) model.jacobian_eval = jac_algebraic elif len(model.algebraic) == 0: # No algebraic equations: residuals is rhs only model.residuals_eval = Residuals(rhs, "residuals", model) model.jacobian_eval = jac_rhs # Calculate consistent initial conditions for the algebraic equations else: all_states = pybamm.NumpyConcatenation( model.concatenated_rhs, model.concatenated_algebraic) # Process again, uses caching so should be quick residuals_eval, jacobian_eval = process(all_states, "residuals")[1:] model.residuals_eval = residuals_eval model.jacobian_eval = jacobian_eval pybamm.logger.info("Finish solver set-up")
def test_evaluator_jax(self): a = pybamm.StateVector(slice(0, 1)) b = pybamm.StateVector(slice(1, 2)) y_tests = [ np.array([[2.0], [3.0]]), np.array([[1.0], [3.0]]), np.array([1.0, 3.0]), ] t_tests = [1.0, 2.0] # test a * b expr = a * b evaluator = pybamm.EvaluatorJax(expr) result = evaluator.evaluate(t=None, y=np.array([[2], [3]])) self.assertEqual(result, 6) result = evaluator.evaluate(t=None, y=np.array([[1], [3]])) self.assertEqual(result, 3) # test function(a*b) expr = pybamm.Function(test_function, a * b) evaluator = pybamm.EvaluatorJax(expr) result = evaluator.evaluate(t=None, y=np.array([[2], [3]])) self.assertEqual(result, 12) # test exp expr = pybamm.exp(a * b) evaluator = pybamm.EvaluatorJax(expr) result = evaluator.evaluate(t=None, y=np.array([[2], [3]])) self.assertEqual(result, np.exp(6)) # test a constant expression expr = pybamm.Scalar(2) * pybamm.Scalar(3) evaluator = pybamm.EvaluatorJax(expr) result = evaluator.evaluate() self.assertEqual(result, 6) # test a larger expression expr = a * b + b + a**2 / b + 2 * a + b / 2 + 4 evaluator = pybamm.EvaluatorJax(expr) for y in y_tests: result = evaluator.evaluate(t=None, y=y) np.testing.assert_allclose(result, expr.evaluate(t=None, y=y)) # test something with time expr = a * pybamm.t evaluator = pybamm.EvaluatorJax(expr) for t, y in zip(t_tests, y_tests): result = evaluator.evaluate(t=t, y=y) self.assertEqual(result, expr.evaluate(t=t, y=y)) # test something with a matrix multiplication A = pybamm.Matrix(np.array([[1, 2], [3, 4]])) expr = A @ pybamm.StateVector(slice(0, 2)) evaluator = pybamm.EvaluatorJax(expr) for t, y in zip(t_tests, y_tests): result = evaluator.evaluate(t=t, y=y) np.testing.assert_allclose(result, expr.evaluate(t=t, y=y)) # test something with a heaviside a = pybamm.Vector(np.array([1, 2])) expr = a <= pybamm.StateVector(slice(0, 2)) evaluator = pybamm.EvaluatorJax(expr) for t, y in zip(t_tests, y_tests): result = evaluator.evaluate(t=t, y=y) np.testing.assert_allclose(result, expr.evaluate(t=t, y=y)) expr = a > pybamm.StateVector(slice(0, 2)) evaluator = pybamm.EvaluatorJax(expr) for t, y in zip(t_tests, y_tests): result = evaluator.evaluate(t=t, y=y) np.testing.assert_allclose(result, expr.evaluate(t=t, y=y)) # test something with a minimum or maximum a = pybamm.Vector(np.array([1, 2])) expr = pybamm.minimum(a, pybamm.StateVector(slice(0, 2))) evaluator = pybamm.EvaluatorJax(expr) for t, y in zip(t_tests, y_tests): result = evaluator.evaluate(t=t, y=y) np.testing.assert_allclose(result, expr.evaluate(t=t, y=y)) expr = pybamm.maximum(a, pybamm.StateVector(slice(0, 2))) evaluator = pybamm.EvaluatorJax(expr) for t, y in zip(t_tests, y_tests): result = evaluator.evaluate(t=t, y=y) np.testing.assert_allclose(result, expr.evaluate(t=t, y=y)) # test something with an index expr = pybamm.Index(A @ pybamm.StateVector(slice(0, 2)), 0) evaluator = pybamm.EvaluatorJax(expr) for t, y in zip(t_tests, y_tests): result = evaluator.evaluate(t=t, y=y) self.assertEqual(result, expr.evaluate(t=t, y=y)) # test something with a sparse matrix-vector multiplication A = pybamm.Matrix(np.array([[1, 2], [3, 4]])) B = pybamm.Matrix(scipy.sparse.csr_matrix(np.array([[1, 0], [0, 4]]))) C = pybamm.Matrix(scipy.sparse.coo_matrix(np.array([[1, 0], [0, 4]]))) expr = A @ B @ C @ pybamm.StateVector(slice(0, 2)) evaluator = pybamm.EvaluatorJax(expr) for t, y in zip(t_tests, y_tests): result = evaluator.evaluate(t=t, y=y) np.testing.assert_allclose(result, expr.evaluate(t=t, y=y)) # test the sparse-scalar multiplication A = pybamm.Matrix(scipy.sparse.csr_matrix(np.array([[1, 0], [0, 4]]))) for expr in [ A * pybamm.t @ pybamm.StateVector(slice(0, 2)), pybamm.t * A @ pybamm.StateVector(slice(0, 2)), ]: evaluator = pybamm.EvaluatorJax(expr) for t, y in zip(t_tests, y_tests): result = evaluator.evaluate(t=t, y=y) np.testing.assert_allclose(result, expr.evaluate(t=t, y=y)) # test the sparse-scalar division A = pybamm.Matrix(scipy.sparse.csr_matrix(np.array([[1, 0], [0, 4]]))) expr = A / (1.0 + pybamm.t) @ pybamm.StateVector(slice(0, 2)) evaluator = pybamm.EvaluatorJax(expr) for t, y in zip(t_tests, y_tests): result = evaluator.evaluate(t=t, y=y) np.testing.assert_allclose(result, expr.evaluate(t=t, y=y)) # test sparse stack A = pybamm.Matrix(scipy.sparse.csr_matrix(np.array([[1, 0], [0, 4]]))) B = pybamm.Matrix(scipy.sparse.csr_matrix(np.array([[2, 0], [5, 0]]))) a = pybamm.StateVector(slice(0, 1)) expr = pybamm.SparseStack(A, a * B) with self.assertRaises(NotImplementedError): evaluator = pybamm.EvaluatorJax(expr) # test sparse mat-mat mult A = pybamm.Matrix(scipy.sparse.csr_matrix(np.array([[1, 0], [0, 4]]))) B = pybamm.Matrix(scipy.sparse.csr_matrix(np.array([[2, 0], [5, 0]]))) a = pybamm.StateVector(slice(0, 1)) expr = A @ (a * B) with self.assertRaises(NotImplementedError): evaluator = pybamm.EvaluatorJax(expr) # test numpy concatenation a = pybamm.Vector(np.array([[1], [2]])) b = pybamm.Vector(np.array([[3]])) expr = pybamm.NumpyConcatenation(a, b) evaluator = pybamm.EvaluatorJax(expr) for t, y in zip(t_tests, y_tests): result = evaluator.evaluate(t=t, y=y) np.testing.assert_allclose(result, expr.evaluate(t=t, y=y)) # test Inner A = pybamm.Matrix(scipy.sparse.csr_matrix(np.array([[1]]))) v = pybamm.StateVector(slice(0, 1)) for expr in [ pybamm.Inner(A, v) @ v, pybamm.Inner(v, A) @ v, pybamm.Inner(v, v) @ v ]: evaluator = pybamm.EvaluatorJax(expr) for t, y in zip(t_tests, y_tests): result = evaluator.evaluate(t=t, y=y) np.testing.assert_allclose(result, expr.evaluate(t=t, y=y))
def add_ghost_nodes(self, symbol, discretised_symbol, bcs): """ Add ghost nodes to a symbol. For Dirichlet bcs, for a boundary condition "y = a at the left-hand boundary", we concatenate a ghost node to the start of the vector y with value "2*a - y1" where y1 is the value of the first node. Similarly for the right-hand boundary condition. For Dirichlet bcs, for a boundary condition "y = a at the left-hand boundary", we concatenate a ghost node to the start of the vector y with value "2*a - y1" where y1 is the value of the first node. Similarly for the right-hand boundary condition. For Neumann bcs, for a boundary condition "dy/dx = b at the left-hand boundary", we concatenate a ghost node to the start of the vector y with value "b*h + y1" where y1 is the value of the first node and h is the mesh size. Similarly for the right-hand boundary condition. Parameters ---------- domain : list of strings The domain of the symbol for which to add ghost nodes bcs : dict of tuples (:class:`pybamm.Scalar`, str) Dictionary (with keys "left" and "right") of boundary conditions. Each boundary condition consists of a value and a flag indicating its type (e.g. "Dirichlet") Returns ------- :class:`pybamm.Symbol` (shape (n+2, n)) `Matrix @ discretised_symbol + bcs_vector`. When evaluated, this gives the discretised_symbol, with appropriate ghost nodes concatenated at each end. """ # get relevant grid points submesh_list = self.mesh.combine_submeshes(*symbol.domain) # Prepare sizes and empty bcs_vector n = submesh_list[0].npts sec_pts = len(submesh_list) bcs_vector = pybamm.Vector(np.array([])) # starts empty lbc_value, lbc_type = bcs["left"] rbc_value, rbc_type = bcs["right"] for i in range(sec_pts): if lbc_value.evaluates_to_number(): lbc_i = lbc_value else: lbc_i = lbc_value[i] if rbc_value.evaluates_to_number(): rbc_i = rbc_value else: rbc_i = rbc_value[i] if lbc_type == "Dirichlet": left_ghost_constant = 2 * lbc_i elif lbc_type == "Neumann": dx = 2 * (submesh_list[0].nodes[0] - submesh_list[0].edges[0]) left_ghost_constant = -dx * lbc_i else: raise ValueError( "boundary condition must be Dirichlet or Neumann, not '{}'".format( lbc_type ) ) if rbc_type == "Dirichlet": right_ghost_constant = 2 * rbc_i elif rbc_type == "Neumann": dx = 2 * (submesh_list[0].edges[-1] - submesh_list[0].nodes[-1]) right_ghost_constant = dx * rbc_i else: raise ValueError( "boundary condition must be Dirichlet or Neumann, not '{}'".format( rbc_type ) ) # concatenate bcs_vector = pybamm.NumpyConcatenation( bcs_vector, left_ghost_constant, pybamm.Vector(np.zeros(n)), right_ghost_constant, ) # Make matrix to calculate ghost nodes bc_factors = {"Dirichlet": -1, "Neumann": 1} left_factor = bc_factors[lbc_type] right_factor = bc_factors[rbc_type] # coo_matrix takes inputs (data, (row, col)) and puts data[i] at the point # (row[i], col[i]) for each index of data. left_ghost_vector = coo_matrix(([left_factor], ([0], [0])), shape=(1, n)) right_ghost_vector = coo_matrix(([right_factor], ([0], [n - 1])), shape=(1, n)) sub_matrix = vstack([left_ghost_vector, eye(n), right_ghost_vector]) # repeat matrix for secondary dimensions # Convert to csr_matrix so that we can take the index (row-slicing), which is # not supported by the default kron format # Note that this makes column-slicing inefficient, but this should not be an # issue matrix = csr_matrix(kron(eye(sec_pts), sub_matrix)) return pybamm.Matrix(matrix) @ discretised_symbol + bcs_vector
def concatenate(self, *symbols): return pybamm.NumpyConcatenation(*symbols)