diff --git a/.github/workflows/cookbook-practice.yml b/.github/workflows/cookbook-practice.yml index 6ebaf0c..a173140 100644 --- a/.github/workflows/cookbook-practice.yml +++ b/.github/workflows/cookbook-practice.yml @@ -1,6 +1,11 @@ name: Run tests in cookbook -on: [push] +on: + push: + pull_request: + types: [opened, synchronize, reopened, closed] + branches: + - '*' jobs: tests_in_cookbook: diff --git a/.github/workflows/doc-deploy.yml b/.github/workflows/doc-deploy.yml index e4b5601..5b0d77f 100644 --- a/.github/workflows/doc-deploy.yml +++ b/.github/workflows/doc-deploy.yml @@ -11,7 +11,7 @@ on: jobs: build_and_deploy_job: - if: github.event_name == 'push' || (github.event_name == 'pull_request' && github.event.action != 'closed') + if: github.repository == 'smallbunnies/Solverz' &&(github.event_name == 'push' || (github.event_name == 'pull_request' && github.event.action != 'closed')) runs-on: ubuntu-latest name: Build and Deploy Job steps: diff --git a/.github/workflows/run-tests.yml b/.github/workflows/run-tests.yml index c0b580e..60488e1 100644 --- a/.github/workflows/run-tests.yml +++ b/.github/workflows/run-tests.yml @@ -1,6 +1,11 @@ name: Run tests -on: [push] +on: + push: + pull_request: + types: [opened, synchronize, reopened, closed] + branches: + - '*' jobs: built_in_tests: diff --git a/Solverz/__init__.py b/Solverz/__init__.py index 1d9d2ee..70b1125 100644 --- a/Solverz/__init__.py +++ b/Solverz/__init__.py @@ -2,7 +2,8 @@ from Solverz.equation.equations import AE, FDAE, DAE from Solverz.equation.param import Param, IdxParam, TimeSeriesParam from Solverz.sym_algebra.symbols import idx, Para, iVar, iAliasVar -from Solverz.sym_algebra.functions import Sign, Abs, transpose, exp, Diag, Mat_Mul, sin, cos, Min, AntiWindUp, Saturation +from Solverz.sym_algebra.functions import (Sign, Abs, transpose, exp, Diag, Mat_Mul, sin, cos, Min, AntiWindUp, + Saturation, heaviside, ln) from Solverz.num_api.custom_function import minmod_flag, minmod from Solverz.variable.variables import Vars, TimeVars, as_Vars from Solverz.solvers import * diff --git a/Solverz/code_printer/make_module.py b/Solverz/code_printer/make_module.py index 9bb34d4..44d24b4 100644 --- a/Solverz/code_printer/make_module.py +++ b/Solverz/code_printer/make_module.py @@ -11,6 +11,7 @@ def __init__(self, variables: Vars | List[Vars], name: str, lang='python', + make_hvp=False, directory=None, jit=False): self.name = name @@ -20,6 +21,7 @@ def __init__(self, self.variables = [variables] else: self.variables = variables + self.make_hvp = make_hvp self.directory = directory self.jit = jit @@ -29,6 +31,7 @@ def render(self): *self.variables, name=self.name, directory=self.directory, - numba=self.jit) + numba=self.jit, + make_hvp=self.make_hvp) else: - raise NotImplemented(f"{self.lang} module renderer not implemented!") \ No newline at end of file + raise NotImplemented(f"{self.lang} module renderer not implemented!") diff --git a/Solverz/code_printer/python/inline/inline_printer.py b/Solverz/code_printer/python/inline/inline_printer.py index a9b977a..0fd2663 100644 --- a/Solverz/code_printer/python/inline/inline_printer.py +++ b/Solverz/code_printer/python/inline/inline_printer.py @@ -52,7 +52,8 @@ def print_J(eqs_type: str, body.extend(param_decla) body.extend(print_trigger(PARAM)) if not sparse: - body.append(Assignment(temp, zeros(EqnAddr.total_size, var_addr.total_size))) + body.append(Assignment(temp, zeros( + EqnAddr.total_size, var_addr.total_size))) body.extend(print_J_blocks(jac, False)) body.append(Return(temp)) else: @@ -67,7 +68,7 @@ def print_J(eqs_type: str, def print_J_blocks(jac: Jac, sparse: bool): eqn_declaration = [] - for eqn_name, jbs_row in jac.blocks.items(): + for eqn_name, jbs_row in jac.blocks_sorted.items(): for var, jb in jbs_row.items(): eqn_declaration.extend(print_J_block(jb, sparse)) @@ -82,17 +83,55 @@ def print_J_block(jb: JacBlock, sparse: bool) -> List: # extend(iVar('row', internal_use=True), eqn_address), # extend(iVar('col', internal_use=True), var_address), # extend(iVar('data', internal_use=True), data)] - raise NotImplementedError("Matrix parameters in sparse Jac not implemented yet!") + raise NotImplementedError( + "Matrix parameters in sparse Jac not implemented yet!") case 'vector' | 'scalar': return [extend(iVar('row', internal_use=True), SolList(*jb.SpEqnAddr.tolist())), - extend(iVar('col', internal_use=True), SolList(*jb.SpVarAddr.tolist())), + extend(iVar('col', internal_use=True), + SolList(*jb.SpVarAddr.tolist())), extend(iVar('data', internal_use=True), jb.SpDeriExpr)] else: return [AddAugmentedAssignment(iVar('J_', internal_use=True)[jb.DenEqnAddr, jb.DenVarAddr], jb.DenDeriExpr)] -def made_numerical(eqs: SymEquations, y: Vars, sparse=False, output_code=False): +def print_Hvp(eqs_type: str, + hvp: Hvp, + EqnAddr: Address, + var_addr: Address, + PARAM: Dict[str, ParamBase], + nstep: int = 0, + sparse=True): + fp = print_Hvp_prototype(eqs_type, + nstep=nstep) + # initialize temp + temp = iVar('Hvp_', internal_use=True) + body = list() + body.extend(print_var(var_addr, + nstep)[0]) + param_decla = print_param(PARAM)[0] + body.extend(param_decla) + body.extend(print_trigger(PARAM)) + if not sparse: + body.append(Assignment(temp, zeros( + EqnAddr.total_size, var_addr.total_size))) + body.extend(print_J_blocks(hvp, False)) + body.append(Return(temp)) + else: + body.extend([Assignment(iVar('row', internal_use=True), SolList()), + Assignment(iVar('col', internal_use=True), SolList()), + Assignment(iVar('data', internal_use=True), SolList())]) + body.extend(print_J_blocks(hvp, True)) + body.append(Return(coo_2_csc(EqnAddr.total_size, var_addr.total_size))) + Jd = FunctionDefinition.from_FunctionPrototype(fp, body) + return pycode(Jd, fully_qualified_modules=False) + + +def made_numerical(eqs: SymEquations, + y: Vars, + sparse=False, + output_code=False, + make_hvp=False): """ factory method of numerical equations """ @@ -111,11 +150,24 @@ def made_numerical(eqs: SymEquations, y: Vars, sparse=False, output_code=False): eqs.PARAM, eqs.nstep, sparse) + code = {'F': code_F, 'J': code_J} + if make_hvp: + eqs.hvp = Hvp(eqs.jac) + code_HVP = print_Hvp(eqs.__class__.__name__, + eqs.hvp, + eqs.a, + eqs.var_address, + eqs.PARAM, + eqs.nstep, + sparse) + code['HVP'] = code_HVP custom_func = dict() custom_func.update(numerical_interface) custom_func.update(parse_trigger_func(eqs.PARAM)) F = Solverzlambdify(code_F, 'F_', modules=[custom_func, 'numpy']) J = Solverzlambdify(code_J, 'J_', modules=[custom_func, 'numpy']) + if make_hvp: + HVP = Solverzlambdify(code_HVP, 'Hvp_', modules=[custom_func, 'numpy']) p = parse_p(eqs.PARAM) print('Complete!') if isinstance(eqs, SymAE) and not isinstance(eqs, SymFDAE): @@ -126,8 +178,10 @@ def made_numerical(eqs: SymEquations, y: Vars, sparse=False, output_code=False): num_eqn = nDAE(eqs.M, F, J, p) else: raise ValueError(f'Unknown equation type {type(eqs)}') + if make_hvp: + num_eqn.HVP = HVP if output_code: - return num_eqn, {'F': code_F, 'J': code_J} + return num_eqn, code else: return num_eqn @@ -216,7 +270,8 @@ class tolist(Function): @classmethod def eval(cls, *args): if len(args) != 1: - raise ValueError(f"Solverz' tolist function accepts only one input.") + raise ValueError( + f"Solverz' tolist function accepts only one input.") def _numpycode(self, printer, **kwargs): return r'((' + printer._print(self.args[0]) + r').tolist())' diff --git a/Solverz/code_printer/python/inline/tests/test_inline_printer.py b/Solverz/code_printer/python/inline/tests/test_inline_printer.py index 53236d7..b01025d 100644 --- a/Solverz/code_printer/python/inline/tests/test_inline_printer.py +++ b/Solverz/code_printer/python/inline/tests/test_inline_printer.py @@ -15,9 +15,12 @@ from Solverz.sym_algebra.symbols import iVar, idx, Para from Solverz.variable.variables import combine_Vars, as_Vars from Solverz.equation.jac import JacBlock, Ones, Jac -from Solverz.sym_algebra.functions import Diag +from Solverz.equation.hvp import Hvp +from Solverz.sym_algebra.functions import Diag, sin, cos, exp from Solverz.code_printer.python.inline.inline_printer import print_J_block, extend, SolList, print_J_blocks, print_J, \ - print_F, made_numerical + print_F, print_Hvp, made_numerical, Solverzlambdify +from Solverz.utilities.address import Address +from Solverz.num_api.custom_function import numerical_interface # %% row = iVar('row', internal_use=True) @@ -41,7 +44,8 @@ def test_jb_printer_scalar_var_scalar_deri(): assert symJb[1] == extend(col, SolList(1, 1, 1)) assert symJb[2] == extend(data, iVar('y') * Ones(3)) symJb = print_J_block(jb, False) - assert symJb[0] == AddAugmentedAssignment(J_[0:3, 1:2], iVar('y') * Ones(3)) + assert symJb[0] == AddAugmentedAssignment( + J_[0:3, 1], iVar('y') * Ones(3)) def test_jb_printer_vector_var_vector_deri(): @@ -57,7 +61,8 @@ def test_jb_printer_vector_var_vector_deri(): assert symJb[1] == extend(col, SolList(*np.arange(0, 9).tolist())) assert symJb[2] == extend(data, iVar('y')) symJb = print_J_block(jb, False) - assert symJb[0] == AddAugmentedAssignment(iVar('J_', internal_use=True)[1:10, 0:9], Diag(iVar('y'))) + assert symJb[0] == AddAugmentedAssignment( + iVar('J_', internal_use=True)[1:10, 0:9], Diag(iVar('y'))) def test_jbs_printer(): @@ -91,7 +96,7 @@ def test_jbs_printer(): assert symJbs[4] == extend(col, SolList(7, 8, 9, 10, 11, 12, 13, 14, 15)) assert symJbs[5] == extend(data, y ** 2) symJbs = print_J_blocks(jac, False) - assert symJbs[0] == AddAugmentedAssignment(J_[0:3, 1:2], y * Ones(3)) + assert symJbs[0] == AddAugmentedAssignment(J_[0:3, 1], y * Ones(3)) assert symJbs[1] == AddAugmentedAssignment(J_[3:12, 7:16], Diag(y ** 2)) @@ -100,7 +105,7 @@ def test_jbs_printer(): v = y_[1:2] g = p_["g"] J_ = zeros((2, 2)) - J_[0:1,1:2] += ones(1) + J_[0:1,1] += ones(1) return J_ """.strip() @@ -173,8 +178,8 @@ def test_print_F_J(): J_ = zeros((6, 6)) J_[0:2,1:3] += diagflat(ones(2)) J_[2:4,0:2] += diagflat(-ones(2)) - J_[4:5,2:3] += -ones(1) - J_[5:6,2:3] += ones(1) + J_[4:5,2] += -ones(1) + J_[5:6,2] += ones(1) return J_ """.strip() @@ -263,5 +268,73 @@ def test_made_numerical(): nF, code = made_numerical(F, y, sparse=True, output_code=True) F0 = nF.F(y, nF.p) J0 = nF.J(y, nF.p) - np.testing.assert_allclose(F0, np.array([2 * 1 + 1, 1 + np.sin(1)]), rtol=1e-8) - np.testing.assert_allclose(J0.toarray(), np.array([[2, 1], [2, 0.54030231]]), rtol=1e-8) + np.testing.assert_allclose(F0, np.array( + [2 * 1 + 1, 1 + np.sin(1)]), rtol=1e-8) + np.testing.assert_allclose(J0.toarray(), np.array( + [[2, 1], [2, 0.54030231]]), rtol=1e-8) + + +def test_hvp_printer(): + jac = Jac() + x = iVar("x") + jac.add_block( + "a", + x[0], + JacBlock( + "a", + slice(0, 1), + x[0], + np.array([1]), + slice(0, 2), + exp(x[0]), + np.array([2.71828183]), + ), + ) + jac.add_block( + "a", + x[1], + JacBlock( + "a", + slice(0, 1), + x[1], + np.array([1]), + slice(0, 2), + cos(x[1]), + np.array([0.54030231]), + ), + ) + jac.add_block( + "b", + x[0], + JacBlock("b", + slice(1, 2), + x[0], + np.ones(1), + slice(0, 2), + 1, + np.array([1])), + ) + jac.add_block( + "b", + x[1], + JacBlock("b", slice(1, 2), x[1], np.ones(1), + slice(0, 2), 2 * x[1], np.array([2])), + ) + + h = Hvp(jac) + eqn_addr = Address() + eqn_addr.add('a', 1) + eqn_addr.add('b', 1) + var_addr = Address() + var_addr.add('x', 2) + + code = print_Hvp('AE', + h, + eqn_addr, + var_addr, + dict()) + + HVP = Solverzlambdify(code, 'Hvp_', [numerical_interface, 'numpy']) + np.testing.assert_allclose(HVP(np.array([1, 2]), dict(), np.array([1, 1])).toarray(), + np.array([[2.71828183, -0.90929743], + [0., 2.]])) diff --git a/Solverz/code_printer/python/module/module_generator.py b/Solverz/code_printer/python/module/module_generator.py index 71654a6..240de04 100644 --- a/Solverz/code_printer/python/module/module_generator.py +++ b/Solverz/code_printer/python/module/module_generator.py @@ -9,40 +9,64 @@ from Solverz.utilities.io import save from Solverz.code_printer.python.utilities import parse_p, parse_trigger_func from Solverz.code_printer.python.module.module_printer import print_F, print_inner_F, print_sub_inner_F, \ - print_J, print_inner_J + print_J, print_inner_J, print_Hvp, print_inner_Hvp from Solverz.equation.equations import Equations as SymEquations from Solverz.num_api.custom_function import numerical_interface from Solverz.variable.variables import Vars, combine_Vars +from Solverz.equation.hvp import Hvp -def render_modules(eqs: SymEquations, y, name, directory=None, numba=False): +def render_modules(eqs: SymEquations, + y, + name, + directory=None, + numba=False, + make_hvp=False): """ factory method of numerical equations """ print(f"Printing python codes of {eqs.name}...") eqs.FormJac(y) p = parse_p(eqs.PARAM) - code_F = print_F(eqs.__class__.__name__, - eqs.var_address, - eqs.PARAM, - eqs.nstep) - code_inner_F = print_inner_F(eqs.EQNs, - eqs.a, - eqs.var_address, - eqs.PARAM, - eqs.nstep) - code_sub_inner_F = print_sub_inner_F(eqs.EQNs) - code_J = print_J(eqs.__class__.__name__, - eqs.eqn_size, - eqs.var_address, - eqs.PARAM, - eqs.nstep) - inner_J = print_inner_J(eqs.var_address, - eqs.PARAM, - eqs.jac, - eqs.nstep) - code_inner_J = inner_J['code_inner_J'] - code_sub_inner_J = inner_J['code_sub_inner_J'] + + code_dict = dict() + code_dict['F'] = print_F(eqs.__class__.__name__, + eqs.var_address, + eqs.PARAM, + eqs.nstep) + code_dict["inner_F"] = print_inner_F(eqs.EQNs, + eqs.a, + eqs.var_address, + eqs.PARAM, + eqs.nstep) + code_dict["sub_inner_F"] = print_sub_inner_F(eqs.EQNs) + + code_dict["J"] = print_J(eqs.__class__.__name__, + eqs.eqn_size, + eqs.var_address, + eqs.PARAM, + eqs.nstep) + J = print_inner_J(eqs.var_address, + eqs.PARAM, + eqs.jac, + eqs.nstep) + code_dict["inner_J"] = J['code_inner_J'] + code_dict["sub_inner_J"] = J['code_sub_inner_J'] + + if make_hvp: + eqs.hvp = Hvp(eqs.jac) + code_dict["Hvp"] = print_Hvp(eqs.__class__.__name__, + eqs.eqn_size, + eqs.var_address, + eqs.PARAM, + eqs.nstep) + inner_Hvp = print_inner_Hvp(eqs.var_address, + eqs.PARAM, + eqs.hvp, + eqs.nstep) + code_dict["inner_Hvp"] = inner_Hvp['code_inner_Hvp'] + code_dict["sub_inner_Hvp"] = inner_Hvp['code_sub_inner_Hvp'] + custom_func = dict() custom_func.update(numerical_interface) @@ -57,10 +81,11 @@ def print_trigger_func_code(): code_tfuc[func_name] = func_code return code_tfuc - code_tfunc_dict = print_trigger_func_code() + code_dict["tfunc_dict"] = print_trigger_func_code() print('Complete!') + # parse equation parameters, which is different the p dict eqn_parameter = {} if isinstance(eqs, SymAE) and not isinstance(eqs, SymFDAE): eqn_type = 'AE' @@ -73,15 +98,18 @@ def print_trigger_func_code(): else: raise ValueError(f'Unknown equation type {type(eqs)}') - code_dict = {'F': code_F, - 'inner_F': code_inner_F, - 'sub_inner_F': code_sub_inner_F, - 'J': code_J, - 'inner_J': code_inner_J, - 'sub_inner_J': code_sub_inner_J, - 'tfunc_dict': code_tfunc_dict} row, col, data = eqs.jac.parse_row_col_data() eqn_parameter.update({'row': row, 'col': col, 'data': data}) + if make_hvp: + row_hvp, col_hvp, data_hvp = eqs.hvp.jac1.parse_row_col_data() + else: + row_hvp = 0 + col_hvp = 0 + data_hvp = 0 + eqn_parameter.update({'row_hvp': row_hvp, + 'col_hvp': col_hvp, + 'data_hvp': data_hvp}) + print(f"Rendering python modules!") render_as_modules(name, code_dict, @@ -100,7 +128,12 @@ def is_valid_python_module_name(module_name): return bool(pattern.match(module_name)) -def create_python_module(module_name, initiate_code, module_code, dependency_code, auxiliary, directory=None): +def create_python_module(module_name, + initiate_code, + module_code, + dependency_code, + auxiliary, + directory=None): location = module_name if directory is None else directory + '/' + module_name # Create the parent directory if it doesn't exist @@ -124,6 +157,25 @@ def create_python_module(module_name, initiate_code, module_code, dependency_cod save(auxiliary, f'{location}/param_and_setting.pkl') +try_hvp = """ +try: + from .num_func import Hvp_ + mdl.HVP = Hvp_ + has_hvp = True +except ImportError: + has_hvp = False +""" + +code_compile = """ +mdl.F({alpha}) +mdl.J({alpha}) +if has_hvp: + from numpy import ones_like + v = ones_like(y) + mdl.HVP({beta}) +""" + + def print_init_code(eqn_type: str, module_name, eqn_param): code = 'from .num_func import F_, J_\n' code += 'from .dependency import setting, p_, y\n' @@ -132,7 +184,8 @@ def print_init_code(eqn_type: str, module_name, eqn_param): case 'AE': code += 'from Solverz.num_api.num_eqn import nAE\n' code += 'mdl = nAE(F_, J_, p_)\n' - code_compile = 'mdl.F(y, p_)\nmdl.J(y, p_)\n' + code_compile_args_F_J = "y, p_" + code_compile_args_Hvp = "y, p_, v" case 'FDAE': try: nstep = eqn_param['nstep'] @@ -141,16 +194,19 @@ def print_init_code(eqn_type: str, module_name, eqn_param): code += 'from Solverz.num_api.num_eqn import nFDAE\n' code += 'mdl = nFDAE(F_, J_, p_, setting["nstep"])\n' args_str = ', '.join(['y' for i in range(nstep)]) - code_compile = f'mdl.F(0, y, p_, {args_str})\nmdl.J(0, y, p_, {args_str})\n' + code_compile_args_F_J = f"0, y, p_, {args_str}" + code_compile_args_Hvp = f"0, y, p_, v, {args_str}" case 'DAE': code += 'from Solverz.num_api.num_eqn import nDAE\n' code += 'mdl = nDAE(setting["M"], F_, J_, p_)\n' - code_compile = 'mdl.F(0, y, p_)\nmdl.J(0, y, p_)\n' + code_compile_args_F_J = "0, y, p_" + code_compile_args_Hvp = "0, y, p_, v" case _: raise ValueError(f'Unknown equation type {eqn_type}') + code += try_hvp code += f'print("Compiling model {module_name}...")\n' code += f'start = time.perf_counter()\n' - code += code_compile + code += code_compile.format(alpha=code_compile_args_F_J, beta=code_compile_args_Hvp) code += f'end = time.perf_counter()\n' code += "print(f'Compiling time elapsed: {end-start}s')\n" return code @@ -159,6 +215,7 @@ def print_init_code(eqn_type: str, module_name, eqn_param): def print_module_code(code_dict: Dict[str, str], numba=False): code = 'from .dependency import *\n' code += """_data_ = setting["data"]\n""" + code += """_data_hvp = setting["data_hvp"]\n""" code += """_F_ = zeros_like(y, dtype=float64)""" code += '\n\r\n' for tfunc in code_dict['tfunc_dict'].values(): @@ -166,30 +223,46 @@ def print_module_code(code_dict: Dict[str, str], numba=False): code += '@njit(cache=True)\n' code += tfunc code += '\n\r\n' + code += code_dict['F'] code += '\n\r\n' if numba: code += '@njit(cache=True)\n' code += code_dict['inner_F'] code += '\n\r\n' - sub_inner_F = code_dict['sub_inner_F'] - for sub_func in sub_inner_F: + for sub_func in code_dict['sub_inner_F']: if numba: code += '@njit(cache=True)\n' code += sub_func code += '\n\r\n' + code += code_dict['J'] code += '\n\r\n' if numba: code += '@njit(cache=True)\n' code += code_dict['inner_J'] code += '\n\r\n' - sub_inner_J = code_dict['sub_inner_J'] - for sub_func in sub_inner_J: + for sub_func in code_dict['sub_inner_J']: if numba: code += '@njit(cache=True)\n' code += sub_func code += '\n\r\n' + + try: + code += code_dict['Hvp'] + code += '\n\r\n' + if numba: + code += '@njit(cache=True)\n' + code += code_dict['inner_Hvp'] + code += '\n\r\n' + for sub_func in code_dict['sub_inner_Hvp']: + if numba: + code += '@njit(cache=True)\n' + code += sub_func + code += '\n\r\n' + except KeyError: + pass + return code @@ -206,13 +279,23 @@ def print_dependency_code(modules): code += 'setting = auxiliary["eqn_param"]\n' code += 'row = setting["row"]\n' code += 'col = setting["col"]\n' + code += 'row_hvp = setting["row_hvp"]\n' + code += 'col_hvp = setting["col_hvp"]\n' + code += 'data_hvp = setting["data_hvp"]\n' code += 'p_ = auxiliary["p"]\n' code += 'y = auxiliary["vars"]\n' return code -def render_as_modules(name, code_dict: Dict[str, str], eqn_type: str, p: dict, eqn_param: dict, variables: Vars, - modules=None, numba=False, directory=None): +def render_as_modules(name, + code_dict: Dict[str, str], + eqn_type: str, + p: dict, + eqn_param: dict, + variables: Vars, + modules=None, + numba=False, + directory=None): if is_valid_python_module_name(name): module_name = name else: diff --git a/Solverz/code_printer/python/module/module_printer.py b/Solverz/code_printer/python/module/module_printer.py index 98e3e00..9818b9a 100644 --- a/Solverz/code_printer/python/module/module_printer.py +++ b/Solverz/code_printer/python/module/module_printer.py @@ -4,6 +4,72 @@ # %% +def print_Hvp(eqs_type: str, + eqn_size: int, + var_addr: Address, + PARAM: Dict[str, ParamBase], + nstep: int = 0): + if eqn_size != var_addr.total_size: + raise ValueError(f"Hvp matrix, with size ({eqn_size}*{var_addr.total_size}), not square") + fp = print_Hvp_prototype(eqs_type, + 'Hvp_', + nstep=nstep) + body = [] + var_assignments, var_list = print_var(var_addr, + nstep) + body.extend(var_assignments) + param_assignments, param_list = print_param(PARAM) + body.extend(param_assignments) + body.extend(print_trigger(PARAM)) + args = [symbols('_data_hvp', real=True), symbols('v_', real=True)] + var_list + param_list + body.extend([Assignment(iVar('data_hvp', internal_use=True), FunctionCall('inner_Hvp', args))]) + body.extend([Return(coo_2_csc_hvp(eqn_size, var_addr.total_size))]) + fd = FunctionDefinition.from_FunctionPrototype(fp, body) + return pycode(fd, fully_qualified_modules=False) + + +def print_inner_Hvp(var_addr: Address, + PARAM: Dict[str, ParamBase], + hvp: Hvp, + nstep: int = 0): + var_assignments, var_list = print_var(var_addr, + nstep) + param_assignments, param_list = print_param(PARAM) + args = [] + for var in var_list + param_list: + args.append(symbols(var.name, real=True)) + fp = FunctionPrototype(real, 'inner_Hvp', + [symbols('_data_hvp', real=True), symbols('v_', real=True)] + args) + body = [] + + code_sub_inner_Hvp_blocks = [] + count = 0 + addr_by_ele_0 = 0 + for eqn_name, jbs_row in hvp.blocks_sorted.items(): + for var, jb in jbs_row.items(): + rhs = jb.SpDeriExpr + SymbolsInDeri_ = list(Eqn(f'temp' + eqn_name + var.name, rhs).SYMBOLS.values()) + # add real assumption + SymbolsInDeri = [symbols(arg.name, real=True) for arg in SymbolsInDeri_] + addr_by_ele = slice(addr_by_ele_0, addr_by_ele_0 + jb.SpEleSize) + if not jb.IsDeriNumber: + # _data_[0:1] = inner_Hvp0(v_, t1, x) + body.append(Assignment(iVar('_data_hvp', internal_use=True)[addr_by_ele], + FunctionCall(f'inner_Hvp{int(count)}', SymbolsInDeri))) + + # def inner_Hvp0(v_, t1, x): + # return -t1 * pi * cos(pi * x) + 1 + fp1 = FunctionPrototype(real, f'inner_Hvp{int(count)}', SymbolsInDeri) + body1 = [Return(rhs)] + fd1 = FunctionDefinition.from_FunctionPrototype(fp1, body1) + code_sub_inner_Hvp_blocks.append(pycode(fd1, fully_qualified_modules=False)) + count += 1 + addr_by_ele_0 += jb.SpEleSize + temp = iVar('_data_hvp', internal_use=True) + body.extend([Return(temp)]) + fd = FunctionDefinition.from_FunctionPrototype(fp, body) + return {'code_inner_Hvp': pycode(fd, fully_qualified_modules=False), + 'code_sub_inner_Hvp': code_sub_inner_Hvp_blocks} def print_J(eqs_type: str, @@ -46,7 +112,7 @@ def print_inner_J(var_addr: Address, code_sub_inner_J_blocks = [] count = 0 addr_by_ele_0 = 0 - for eqn_name, jbs_row in jac.blocks.items(): + for eqn_name, jbs_row in jac.blocks_sorted.items(): for var, jb in jbs_row.items(): rhs = jb.SpDeriExpr SymbolsInDeri_ = list(Eqn(f'temp' + eqn_name + var.name, rhs).SYMBOLS.values()) diff --git a/Solverz/code_printer/python/module/test/test_module_generator.py b/Solverz/code_printer/python/module/test/test_module_generator.py index 4e79702..ed965d6 100644 --- a/Solverz/code_printer/python/module/test/test_module_generator.py +++ b/Solverz/code_printer/python/module/test/test_module_generator.py @@ -6,7 +6,8 @@ import numpy as np from sympy import symbols, pycode, Integer -from Solverz import as_Vars, Eqn, Ode, AE, sin, made_numerical, Model, Var, Param, TimeSeriesParam, AliasVar, Abs +from Solverz import as_Vars, Eqn, Ode, AE, sin, made_numerical, Model, Var, Param, TimeSeriesParam, AliasVar, Abs, exp, \ + cos from Solverz.code_printer.make_module import module_printer from Solverz.code_printer.python.inline.inline_printer import print_J_block from Solverz.code_printer.python.utilities import _print_var_parser @@ -24,6 +25,9 @@ setting = auxiliary["eqn_param"] row = setting["row"] col = setting["col"] +row_hvp = setting["row_hvp"] +col_hvp = setting["col_hvp"] +data_hvp = setting["data_hvp"] p_ = auxiliary["p"] y = auxiliary["vars"] """ @@ -33,10 +37,22 @@ import time from Solverz.num_api.num_eqn import nAE mdl = nAE(F_, J_, p_) + +try: + from .num_func import Hvp_ + mdl.HVP = Hvp_ + has_hvp = True +except ImportError: + has_hvp = False print("Compiling model test_eqn1...") start = time.perf_counter() + mdl.F(y, p_) mdl.J(y, p_) +if has_hvp: + from numpy import ones_like + v = ones_like(y) + mdl.HVP(y, p_, v) end = time.perf_counter() print(f'Compiling time elapsed: {end-start}s') """ @@ -116,6 +132,83 @@ def test_AE_module_printer(): shutil.rmtree(test_folder_path) +expected_Hvp_ = """def Hvp_(y_, p_, v_): + x = y_[0:2] + c = p_["c"] + data_hvp = inner_Hvp(_data_hvp, v_, x, c) + return coo_array((data_hvp, (row_hvp, col_hvp)), (2, 2)).tocsc() +""" + +expected_inner_Hvp = """@njit(cache=True) +def inner_Hvp(_data_hvp, v_, x, c): + _data_hvp[0:1] = inner_Hvp0(v_, x) + _data_hvp[1:2] = inner_Hvp1(v_, x) + _data_hvp[2:3] = inner_Hvp2(c, v_) + return _data_hvp +""" + +expected_inner_Hvp0 = """@njit(cache=True) +def inner_Hvp0(v_, x): + return v_[0]*exp(x[0])*ones(1) +""" + +expected_inner_Hvp1 = """@njit(cache=True) +def inner_Hvp1(v_, x): + return -v_[1]*sin(x[1])*ones(1) +""" + +expected_inner_Hvp2 = """@njit(cache=True) +def inner_Hvp2(c, v_): + return v_[1]*c*ones(1) +""" + + +def test_AE_module_generator_with_hvp(): + m = Model() + m.x = Var('x', [1, 1]) + m.c = Param('c', 2) + m.f1 = Eqn('f1', 1 + exp(m.x[0]) + sin(m.x[1])) + m.f2 = Eqn('f2', m.x[0] + m.c / 2 * m.x[1] ** 2) + + F, y0 = m.create_instance() + + current_file_path = os.path.abspath(__file__) + current_folder = os.path.dirname(current_file_path) + + test_folder_path = current_folder + '\\Solverz_test_ae_module_printer_with_hvp' + sys.path.extend([test_folder_path]) + + pyprinter = module_printer(F, + y0, + 'test_module', + jit=True, + directory=test_folder_path, + make_hvp=True) + pyprinter.render() + + from test_module import mdl, y + from test_module.num_func import Hvp_, inner_Hvp, inner_Hvp0, inner_Hvp1, inner_Hvp2 + + F0 = mdl.F(y, mdl.p) + J0 = mdl.J(y, mdl.p) + Hvp0 = mdl.HVP(y, mdl.p, np.ones(2)) + np.testing.assert_allclose(F0, np.array([4.559752813, 2]), rtol=1e-8) + np.testing.assert_allclose(J0.toarray(), np.array([[2.71828183, 0.54030231], + [1., 2.]]), rtol=1e-8) + np.testing.assert_allclose(Hvp0.toarray(), + np.array([[2.71828183, -0.84147098], + [0., 2.]]), + rtol=1e-8) + + assert inspect.getsource(Hvp_) == expected_Hvp_ + assert inspect.getsource(inner_Hvp) == expected_inner_Hvp + assert inspect.getsource(inner_Hvp0.func_code) == expected_inner_Hvp0 + assert inspect.getsource(inner_Hvp1.func_code) == expected_inner_Hvp1 + assert inspect.getsource(inner_Hvp2.func_code) == expected_inner_Hvp2 + + shutil.rmtree(test_folder_path) + + expected_F = """def F_(t, y_, p_, y_0): p = y_[0:82] q = y_[82:164] @@ -148,10 +241,10 @@ def inner_F(_F_, p, q, p_tag_0, q_tag_0, pb, qb): expected_inner_J = """@njit(cache=True) def inner_J(_data_, p, q, p_tag_0, q_tag_0, pb, qb): - _data_[0:81] = inner_J0(p, p_tag_0, q, q_tag_0) - _data_[81:162] = inner_J1(p, p_tag_0, q, q_tag_0) - _data_[162:243] = inner_J2(p, p_tag_0, q, q_tag_0) - _data_[243:324] = inner_J3(p, p_tag_0, q, q_tag_0) + _data_[2:83] = inner_J0(p, p_tag_0, q, q_tag_0) + _data_[83:164] = inner_J1(p, p_tag_0, q, q_tag_0) + _data_[164:245] = inner_J2(p, p_tag_0, q, q_tag_0) + _data_[245:326] = inner_J3(p, p_tag_0, q, q_tag_0) return _data_ """ diff --git a/Solverz/code_printer/python/module/test/test_module_printer.py b/Solverz/code_printer/python/module/test/test_module_printer.py index d6c31b0..ac639ee 100644 --- a/Solverz/code_printer/python/module/test/test_module_printer.py +++ b/Solverz/code_printer/python/module/test/test_module_printer.py @@ -9,10 +9,155 @@ from Solverz.equation.eqn import Eqn, Ode from Solverz.equation.param import Param, TimeSeriesParam from Solverz.equation.jac import Jac, JacBlock +from Solverz.equation.hvp import Hvp +from Solverz.sym_algebra.functions import exp, cos from Solverz.utilities.address import Address from sympy.codegen.ast import FunctionDefinition, Return from Solverz.code_printer.python.module.module_printer import (print_J, print_inner_J, print_F, - print_inner_F, print_sub_inner_F) + print_inner_F, print_sub_inner_F, + print_Hvp, print_inner_Hvp) + +expected_AE_Hvp = """def Hvp_(y_, p_, v_): + omega = y_[0:10] + delta = y_[10:15] + x = y_[15:18] + y = y_[18:25] + ax = p_["ax"] + lam = p_["lam"] + ax = ax_trigger_func(x) + data_hvp = inner_Hvp(_data_hvp, v_, omega, delta, x, y, ax, lam) + return coo_array((data_hvp, (row_hvp, col_hvp)), (25, 25)).tocsc() +""".strip() +expected_FDAE_Hvp = """def Hvp_(t, y_, p_, v_, y_0): + omega = y_[0:10] + delta = y_[10:15] + x = y_[15:18] + y = y_[18:25] + omega_tag_0 = y_0[0:10] + delta_tag_0 = y_0[10:15] + x_tag_0 = y_0[15:18] + y_tag_0 = y_0[18:25] + ax = p_["ax"] + lam = p_["lam"] + ax = ax_trigger_func(x) + data_hvp = inner_Hvp(_data_hvp, v_, omega, delta, x, y, omega_tag_0, delta_tag_0, x_tag_0, y_tag_0, ax, lam) + return coo_array((data_hvp, (row_hvp, col_hvp)), (25, 25)).tocsc() +""".strip() + + +def test_print_Hvp(): + eqs_type = 'AE' + VarAddr = Address() + VarAddr.add('omega', 10) + VarAddr.add('delta', 5) + VarAddr.add('x', 3) + VarAddr.add('y', 7) + Pdict = dict() + ax = Param('ax', + [0, 1], + triggerable=True, + trigger_var=['x'], + trigger_fun=np.sin) + Pdict['ax'] = ax + lam = Param('lam', [1, 1, 1]) + Pdict["lam"] = lam + + with pytest.raises(ValueError, match=re.escape("Hvp matrix, with size (20*25), not square")): + print_Hvp(eqs_type, + 20, + VarAddr, + Pdict) + + assert print_Hvp(eqs_type, + 25, + VarAddr, + Pdict) == expected_AE_Hvp + + assert print_Hvp("FDAE", + 25, + VarAddr, + Pdict, + nstep=1) == expected_FDAE_Hvp + + +expected_inner_Hvp = """def inner_Hvp(_data_hvp, v_, x, c): + _data_hvp[0:1] = inner_Hvp0(v_, x) + _data_hvp[1:2] = inner_Hvp1(v_, x) + _data_hvp[2:3] = inner_Hvp2(c, v_) + return _data_hvp +""".strip() +expected_inner_Hvp0 = """def inner_Hvp0(v_, x): + return v_[0]*exp(x[0])*ones(1) +""".strip() +expected_inner_Hvp1 = """def inner_Hvp1(v_, x): + return -v_[1]*sin(x[1])*ones(1) +""".strip() +expected_inner_Hvp2 = """def inner_Hvp2(c, v_): + return v_[1]*c*ones(1) +""".strip() + + +def test_print_inner_Hvp(): + jac = Jac() + x = iVar("x") + c = Para('c') + jac.add_block( + "a", + x[0], + JacBlock( + "a", + slice(0, 1), + x[0], + np.array([1]), + slice(0, 2), + exp(x[0]), + np.array([2.71828183]), + ), + ) + jac.add_block( + "a", + x[1], + JacBlock( + "a", + slice(0, 1), + x[1], + np.array([1]), + slice(0, 2), + cos(x[1]), + np.array([0.54030231]), + ), + ) + jac.add_block( + "b", + x[0], + JacBlock("b", + slice(1, 2), + x[0], + np.ones(1), + slice(0, 2), + Integer(1), + np.array([1])), + ) + jac.add_block( + "b", + x[1], + JacBlock("b", slice(1, 2), x[1], np.ones(1), slice(0, 2), c * x[1], np.array([2])), + ) + + Pdict = {"c": Param("c", np.array([2.0]))} + VarAddr = Address() + VarAddr.add('x', 2) + + h = Hvp(jac) + code_dict = print_inner_Hvp(VarAddr, + Pdict, + h) + + assert code_dict['code_inner_Hvp'] == expected_inner_Hvp + assert code_dict['code_sub_inner_Hvp'][0] == expected_inner_Hvp0 + assert code_dict['code_sub_inner_Hvp'][1] == expected_inner_Hvp1 + assert code_dict['code_sub_inner_Hvp'][2] == expected_inner_Hvp2 + expected = """def J_(t, y_, p_): omega = y_[0:10] diff --git a/Solverz/code_printer/python/utilities.py b/Solverz/code_printer/python/utilities.py index 7d0649c..8706210 100644 --- a/Solverz/code_printer/python/utilities.py +++ b/Solverz/code_printer/python/utilities.py @@ -16,6 +16,7 @@ from Solverz.equation.eqn import Eqn from Solverz.equation.equations import Equations as SymEquations, FDAE as SymFDAE, DAE as SymDAE, AE as SymAE from Solverz.equation.jac import JacBlock, Jac +from Solverz.equation.hvp import Hvp from Solverz.equation.param import TimeSeriesParam, ParamBase from Solverz.sym_algebra.symbols import iVar, SolDict, Para, idx, IdxSymBasic from Solverz.sym_algebra.functions import Arange @@ -127,9 +128,36 @@ def print_F_J_prototype(eqs_type: str, func_name: str, nstep=0): if eqs_type == 'FDAE' and nstep == 0: raise ValueError("nstep of FDAE should be greater than 0!") elif eqs_type == 'FDAE' and nstep > 0: - xtra_args.extend([symbols('y_' + f'{i}', real=True) for i in range(nstep)]) + xtra_args.extend( + [symbols('y_' + f'{i}', real=True) for i in range(nstep)]) elif eqs_type != 'FDAE' and nstep > 0: - raise TypeError(f"Only FDAE supports the case where nstep > 0! Not {eqs_type}") + raise TypeError( + f"Only FDAE supports the case where nstep > 0! Not {eqs_type}") + else: # eqs_type != 'FDAE' and nstep == 0 + pass + + fp = FunctionPrototype(real, func_name, args + xtra_args) + return fp + + +def print_Hvp_prototype(eqs_type: str, func_name: str = 'Hvp_', nstep=0): + t, y_, p_, v_ = symbols('t y_ p_ v_', real=True) + if eqs_type == 'DAE' or eqs_type == 'FDAE': + args = [t, y_, p_, v_] + elif eqs_type == 'AE': + args = [y_, p_, v_] + else: + raise TypeError(f"Unknown equation type {eqs_type}") + xtra_args = [] + + if eqs_type == 'FDAE' and nstep == 0: + raise ValueError("nstep of FDAE should be greater than 0!") + elif eqs_type == 'FDAE' and nstep > 0: + xtra_args.extend( + [symbols('y_' + f'{i}', real=True) for i in range(nstep)]) + elif eqs_type != 'FDAE' and nstep > 0: + raise TypeError( + f"Only FDAE supports the case where nstep > 0! Not {eqs_type}") else: # eqs_type != 'FDAE' and nstep == 0 pass @@ -153,7 +181,8 @@ def print_eqn_assignment(EQNs: Dict[str, Eqn], for var in list(eqn.expr.free_symbols): if isinstance(var, IdxSymBasic): if isinstance(var.index, (idx, list)): - raise ValueError(f"Numba printer does not accept idx or list index {var.index}") + raise ValueError( + f"Numba printer does not accept idx or list index {var.index}") _F_ = iVar('_F_', internal_use=True) eqn_declaration.append(Assignment(_F_[eqn_address], FunctionCall(f'inner_F{int(count)}', list(eqn.SYMBOLS.values())))) @@ -179,12 +208,28 @@ def _pythoncode(self, printer, **kwargs): return self._numpycode(printer, **kwargs) +class coo_2_csc_hvp(Symbol): + + def __new__(cls, eqn_size: int, vsize: int): + obj = Symbol.__new__(cls, f'coo_2_csc') + obj.eqn_size = eqn_size + obj.vsize = vsize + return obj + + def _numpycode(self, printer, **kwargs): + return f'coo_array((data_hvp, (row_hvp, col_hvp)), ({self.eqn_size}, {self.vsize})).tocsc()' + + def _pythoncode(self, printer, **kwargs): + return self._numpycode(printer, **kwargs) + + class coo_array(Function): @classmethod def eval(cls, *args): if len(args) > 1: - raise ValueError(f"Solverz' coo_array object accepts only one inputs.") + raise ValueError( + f"Solverz' coo_array object accepts only one inputs.") def _numpycode(self, printer, **kwargs): return f'coo_array({printer._print(self.args[0])})' diff --git a/Solverz/equation/equations.py b/Solverz/equation/equations.py index ffa7e7d..0b21f0e 100644 --- a/Solverz/equation/equations.py +++ b/Solverz/equation/equations.py @@ -6,6 +6,7 @@ from copy import deepcopy import numpy as np +from Solverz.equation.hvp import Hvp from sympy import Symbol, Integer, Expr, Number as SymNumber from scipy.sparse import csc_array, coo_array # from cvxopt import spmatrix, matrix @@ -16,6 +17,7 @@ from Solverz.sym_algebra.functions import Slice from Solverz.variable.variables import Vars from Solverz.utilities.address import Address, combine_Address +from Solverz.utilities.type_checker import is_integer from Solverz.num_api.Array import Array from Solverz.equation.jac import Jac, JacBlock @@ -540,25 +542,14 @@ def M(self): elif isinstance(diff_var, IdxVar): var_idx = diff_var.index var_name = diff_var.name0 - if isinstance(var_idx, (np.integer, int, Integer)): + if is_integer(var_idx): variable_address = self.var_address.v[var_name][var_idx: var_idx + 1] elif isinstance(var_idx, str): variable_address = self.var_address.v[var_name][np.ix_(self.PARAM[var_idx].v.reshape((-1,)))] elif isinstance(var_idx, slice): - temp = [] - if var_idx.start is not None: - temp.append(var_idx.start) - if var_idx.stop is not None: - temp.append(var_idx.stop) - if var_idx.step is not None: - temp.append(var_idx.step) - temp_func = Eqn('To evaluate var_idx of variable' + diff_var.name, Slice(*temp)) - variable_address = self.var_address.v[var_name][ - temp_func.NUM_EQN(*self.obtain_eqn_args(temp_func))] + variable_address = self.var_address.v[var_name][var_idx] elif isinstance(var_idx, Expr): - temp_func = Eqn('To evaluate var_idx of variable' + diff_var.name, var_idx) - args = self.obtain_eqn_args(temp_func) - variable_address = self.var_address.v[var_name][temp_func.NUM_EQN(*args).reshape(-1, )] + raise TypeError(f"Index of {diff_var} cannot be sympy.Expr!") elif isinstance(var_idx, list): variable_address = self.var_address.v[var_name][var_idx] else: diff --git a/Solverz/equation/hvp.py b/Solverz/equation/hvp.py new file mode 100644 index 0000000..b29980f --- /dev/null +++ b/Solverz/equation/hvp.py @@ -0,0 +1,83 @@ +from typing import Dict, Union +from sympy import Integer, Expr + +from Solverz.sym_algebra.symbols import iVar, IdxVar, Para +from Solverz.equation.jac import Jac, JacBlock +from Solverz.utilities.address import Address +from Solverz.utilities.type_checker import is_zero + +SolVar = Union[iVar, IdxVar] + + +class Hvp: + + def __init__(self, jac: Jac) -> None: + self.blocks_sorted: Dict[str, Dict[SolVar, JacBlock]] = dict() + self.jac0 = jac + self.jac1 = Jac() + + # first multiply jac by vector v + self.eqn_column: Dict[str, Expr] = dict() + v = Para("v_", internal_use=True) + for eqn_name, jbs_row in self.jac0.blocks.items(): + expr = Integer(0) + for var, jb in jbs_row.items(): + den_var_addr = parse_den_var_addr(jb.DenVarAddr) + match jb.DiffVarType: + case "scalar": + match jb.DeriType: + case "scalar": + expr += jb.DeriExpr * v[den_var_addr] + case "vector": + expr += jb.DeriExpr * v[den_var_addr] + case _: + raise TypeError( + f"Unknown Derivative type {jb.DiffVarType}!" + ) + case "vector": + match jb.DeriType: + case "scalar": + expr += jb.DeriExpr * v[den_var_addr] + case "vector": + expr += jb.DeriExpr * v[den_var_addr] + case "matrix": + raise NotImplementedError( + "Matrix derivative not implemented in Hvp!" + ) + case _: + raise TypeError( + f"Unknown Derivative type {jb.DiffVarType}!" + ) + case _: + raise TypeError(f"Unknown DiffVarType {jb.DiffVarType}!") + self.eqn_column[eqn_name] = expr + + # derive Jac + for eqn_name, jbs_row in self.jac0.blocks.items(): + for var, jb in jbs_row.items(): + DeriExpr = self.eqn_column[eqn_name].diff(jb.DiffVar) + if not is_zero(DeriExpr): + self.jac1.add_block( + eqn_name, + var, + JacBlock( + eqn_name, + jb.EqnAddr, + jb.DiffVar, + jb.DiffVarValue, + jb.VarAddr, + DeriExpr, + jb.Value0, + ), + ) + + self.blocks_sorted = self.jac1.blocks_sorted + + +def parse_den_var_addr(den_var_addr: slice | int): + if isinstance(den_var_addr, int): + den_var_addr = slice(den_var_addr, den_var_addr + 1) + if den_var_addr.stop - den_var_addr.start == 1: + return den_var_addr.start + else: + return den_var_addr diff --git a/Solverz/equation/jac.py b/Solverz/equation/jac.py index 895979b..66878f4 100644 --- a/Solverz/equation/jac.py +++ b/Solverz/equation/jac.py @@ -15,6 +15,8 @@ class Jac: def __init__(self): self.blocks: Dict[str, Dict[SolVar, JacBlock]] = dict() + self.is_sorted = False + self.__blocks_sorted: Dict[str, Dict[SolVar, JacBlock]] = dict() def add_block(self, eqn_name: str, @@ -25,6 +27,26 @@ def add_block(self, else: self.blocks[eqn_name] = dict() self.blocks[eqn_name][diff_var] = jb + self.is_sorted = False + + def sort_blocks(self): + sorted_blocks = dict() + for eqn_name in sorted(self.blocks): + dict_var_jb = self.blocks[eqn_name] + # SolVar cannot be directly used in sorted, so translate it to string first, then recover + dict_varname_var = {var.__repr__(): var for var in dict_var_jb.keys()} + sorted_blocks[eqn_name] = {dict_varname_var[varname]: dict_var_jb[dict_varname_var[varname]] for varname in + sorted(dict_varname_var)} + self.__blocks_sorted = sorted_blocks + self.is_sorted = True + + @property + def blocks_sorted(self) -> Dict[str, Dict[SolVar, JacBlock]]: + if self.is_sorted: + return self.__blocks_sorted + else: + self.sort_blocks() + return self.__blocks_sorted @property def JacEleNum(self) -> int: @@ -41,11 +63,13 @@ def parse_row_col_data(self): """ Parse the row, col and data for sparse coo-jac construction. """ + if not self.is_sorted: + self.sort_blocks() row = np.zeros(self.JacEleNum, dtype=int) col = np.zeros(self.JacEleNum, dtype=int) data = np.zeros(self.JacEleNum, dtype=float) addr_by_ele_0 = 0 - for eqn_name, jbs_row in self.blocks.items(): + for eqn_name, jbs_row in self.blocks_sorted.items(): for var, jb in jbs_row.items(): addr_by_ele = slice(addr_by_ele_0, addr_by_ele_0 + jb.SpEleSize) row[addr_by_ele] = jb.SpEqnAddr.copy() @@ -88,7 +112,7 @@ def __init__(self, self.SpEleSize = 0 self.SpDeriExpr: Expr = Integer(0) self.DenEqnAddr: slice = slice(0) - self.DenVarAddr: slice = slice(0) + self.DenVarAddr: slice | int = slice(0) self.DenDeriExpr: Expr = Integer(0) EqnSize = self.EqnAddr.stop - self.EqnAddr.start @@ -247,21 +271,24 @@ def ParseDen(self): case 'vector' | 'scalar': self.DenEqnAddr = self.EqnAddr if isinstance(self.DiffVar, iVar): - self.DenVarAddr = self.VarAddr + self.DenVarAddr = self.VarAddr.start else: if isinstance(self.DiffVar.index, slice): VarArange = slice2array(self.VarAddr)[self.DiffVar.index] if VarArange.size > 1: raise ValueError(f"Length of scalar variable {self.DiffVar} > 1!") else: - self.DenVarAddr = slice(VarArange[0], VarArange[-1] + 1) + self.DenVarAddr = VarArange[0] elif is_integer(self.DiffVar.index): idx = int(slice2array(self.VarAddr)[self.DiffVar.index]) - self.DenVarAddr = slice(idx, idx + 1) + self.DenVarAddr = idx else: raise TypeError(f"Index type {type(self.DiffVar.index)} not supported!") self.DenDeriExpr = self.DeriExprBc + def __repr__(self): + return f"Jacblock with DeriExpr {self.DeriExpr.__repr__()}" + def slice2array(s: slice) -> np.ndarray: return np.arange(s.start, s.stop) diff --git a/Solverz/equation/test/test_Param.py b/Solverz/equation/test/test_Param.py index 518b887..0b70ded 100644 --- a/Solverz/equation/test/test_Param.py +++ b/Solverz/equation/test/test_Param.py @@ -1,6 +1,8 @@ import numpy as np +from scipy.sparse import coo_array from Solverz.equation.param import Param, IdxParam, TimeSeriesParam +from Solverz.utilities.testing import assert_allclose_sparse def test_Param(): @@ -20,14 +22,14 @@ def test_Param(): assert Ts3.v.__str__() == '[[1.]\n [2.]\n [3.]]' Ts4 = Param(name='Ts', value=[1, 2, 3], dim=2, sparse=True) - assert Ts4.v.__str__() == ' (0, 0)\t1.0\n (1, 0)\t2.0\n (2, 0)\t3.0' + assert_allclose_sparse(Ts4.v, coo_array(np.array([1, 2, 3]).reshape((-1, 1)).astype(float))) Ts5 = Param(name='Ts', value=[1, 2, 3], dim=2, sparse=False) assert Ts5.v.__str__() == '[[1.]\n [2.]\n [3.]]' - A = np.array([[1, 0], [2, 9], [0, 3]]) - A = Param(name='A', value=A, dim=2, sparse=True) - assert A.v.__str__() == ' (0, 0)\t1.0\n (1, 0)\t2.0\n (1, 1)\t9.0\n (2, 1)\t3.0' + v = np.array([[1, 0], [2, 9], [0, 3]]) + A = Param(name='A', value=v, dim=2, sparse=True) + assert_allclose_sparse(A.v, coo_array(v)) # test of IdxParam f = IdxParam(name='f', diff --git a/Solverz/equation/test/test_hvp.py b/Solverz/equation/test/test_hvp.py new file mode 100644 index 0000000..09ab3c4 --- /dev/null +++ b/Solverz/equation/test/test_hvp.py @@ -0,0 +1,71 @@ +from numbers import Number + +import numpy as np +from numpy.testing import assert_allclose +import re +import pytest + +from sympy import Integer + +from Solverz.equation.eqn import Eqn, Ode +from Solverz.equation.equations import DAE, AE +from Solverz.equation.param import Param +from Solverz.sym_algebra.symbols import iVar, idx, Para +from Solverz.sym_algebra.functions import exp, sin, cos, Diag +from Solverz.variable.variables import combine_Vars, as_Vars +from Solverz.equation.jac import JacBlock, Ones, Jac +from Solverz.equation.hvp import Hvp + + +def test_hvp(): + jac = Jac() + x = iVar("x") + jac.add_block( + "a", + x[0], + JacBlock( + "a", + slice(0, 1), + x[0], + np.array([1]), + slice(0, 2), + exp(x[0]), + np.array([2.71828183]), + ), + ) + jac.add_block( + "a", + x[1], + JacBlock( + "a", + slice(0, 1), + x[1], + np.array([1]), + slice(0, 2), + cos(x[1]), + np.array([0.54030231]), + ), + ) + jac.add_block( + "b", + x[0], + JacBlock("b", + slice(1, 2), + x[0], + np.ones(1), + slice(0, 2), + 1, + np.array([1])), + ) + jac.add_block( + "b", + x[1], + JacBlock("b", slice(1, 2), x[1], np.ones(1), slice(0, 2), 2 * x[1], np.array([2])), + ) + + h = Hvp(jac) + v_ = Para("v_", internal_use=True) + assert h.blocks_sorted['a'][x[0]].DeriExpr == v_[0]*exp(x[0]) + # ISSUE: the two symbolic expressions below cannot be equal, we believe this is a sympy issue. + assert h.blocks_sorted['a'][x[1]].DeriExpr.__repr__() == (-v_[1]*sin(x[1])).__repr__() + assert h.blocks_sorted['b'][x[1]].DeriExpr == 2*v_[1] diff --git a/Solverz/equation/test/test_jac.py b/Solverz/equation/test/test_jac.py index 454302f..3768300 100644 --- a/Solverz/equation/test/test_jac.py +++ b/Solverz/equation/test/test_jac.py @@ -71,7 +71,7 @@ def test_jb_scalar_var_scalar_deri(): np.array([1])) assert jb.DenEqnAddr == slice(0, 3) - assert jb.DenVarAddr == slice(1, 2) + assert jb.DenVarAddr == 1 assert jb.DenDeriExpr == iVar('y') * Ones(3) assert_allclose(jb.SpEqnAddr, np.array([0, 1, 2])) assert_allclose(jb.SpVarAddr, np.array([1, 1, 1])) @@ -89,7 +89,7 @@ def test_jb_scalar_var_scalar_deri(): np.array([1])) assert jb.DenEqnAddr == slice(0, 3) - assert jb.DenVarAddr == slice(2, 3) + assert jb.DenVarAddr == 2 assert jb.DenDeriExpr == iVar('y') * Ones(3) assert_allclose(jb.SpEqnAddr, np.array([0, 1, 2])) assert_allclose(jb.SpVarAddr, np.array([2, 2, 2])) @@ -107,7 +107,7 @@ def test_jb_scalar_var_scalar_deri(): np.array([1])) assert jb.DenEqnAddr == slice(0, 3) - assert jb.DenVarAddr == slice(2, 3) + assert jb.DenVarAddr == 2 assert jb.DenDeriExpr == iVar('y') * Ones(3) assert_allclose(jb.SpEqnAddr, np.array([0, 1, 2])) assert_allclose(jb.SpVarAddr, np.array([2, 2, 2])) @@ -128,7 +128,7 @@ def test_jb_scalar_var_vector_deri(): np.array([1, 1, 1])) assert jb.DenEqnAddr == slice(0, 3) - assert jb.DenVarAddr == slice(1, 2) + assert jb.DenVarAddr == 1 assert jb.DenDeriExpr == iVar('y') assert_allclose(jb.SpEqnAddr, np.array([0, 1, 2])) assert_allclose(jb.SpVarAddr, np.array([1, 1, 1])) @@ -146,7 +146,7 @@ def test_jb_scalar_var_vector_deri(): np.array([1, 1, 1])) assert jb.DenEqnAddr == slice(0, 3) - assert jb.DenVarAddr == slice(2, 3) + assert jb.DenVarAddr == 2 assert jb.DenDeriExpr == iVar('y') assert_allclose(jb.SpEqnAddr, np.array([0, 1, 2])) assert_allclose(jb.SpVarAddr, np.array([2, 2, 2])) @@ -164,7 +164,7 @@ def test_jb_scalar_var_vector_deri(): np.array([1, 1, 1])) assert jb.DenEqnAddr == slice(0, 3) - assert jb.DenVarAddr == slice(2, 3) + assert jb.DenVarAddr == 2 assert jb.DenDeriExpr == iVar('y') assert_allclose(jb.SpEqnAddr, np.array([0, 1, 2])) assert_allclose(jb.SpVarAddr, np.array([2, 2, 2])) diff --git a/Solverz/num_api/custom_function.py b/Solverz/num_api/custom_function.py index 68c0982..6cbd9dc 100644 --- a/Solverz/num_api/custom_function.py +++ b/Solverz/num_api/custom_function.py @@ -82,6 +82,12 @@ def minmod_flag(*args): return np.select(conditions, choice_list, 3) +@implements_nfunc('Heaviside') +@njit +def Heaviside(x): + return np.where(x >= 0, 1.0, 0.0) + + @implements_nfunc('switch') @njit def switch(*args): diff --git a/Solverz/num_api/test/test_Array.py b/Solverz/num_api/test/test_Array.py index 6128dfd..153c9a1 100644 --- a/Solverz/num_api/test/test_Array.py +++ b/Solverz/num_api/test/test_Array.py @@ -1,7 +1,7 @@ import numpy as np -from scipy.sparse import csc_array - +from scipy.sparse import csc_array, coo_array from Solverz.num_api.Array import Array +from Solverz.utilities.testing import assert_allclose_sparse def test_Array(): @@ -10,7 +10,8 @@ def test_Array(): assert Array(x1, dim=1, dtype=float).__repr__() == 'array([1.])' assert Array(x1, dim=2, dtype=float).__repr__() == 'array([[1.]])' assert Array(x1, dim=2, sparse=False, dtype=float).__repr__() == 'array([[1.]])' - assert Array(x1, dim=2, sparse=True, dtype=float).__str__() == ' (0, 0)\t1.0' + assert_allclose_sparse(Array(x1, dim=2, sparse=True, dtype=float), + coo_array(np.array(x1).reshape((-1, 1)).astype(float))) # input as list x2 = [1, 2, 3] @@ -21,7 +22,8 @@ def test_Array(): Array(x2, dim=1, sparse=True) except ValueError as e: assert e.args[0] == 'Cannot create sparse matrix with dim: 1' - assert Array(x2, dim=2, sparse=True, dtype=float).__str__() == ' (0, 0)\t1.0\n (1, 0)\t2.0\n (2, 0)\t3.0' + assert_allclose_sparse(Array(x2, dim=2, sparse=True, dtype=float), + coo_array(np.array(x2).reshape((-1, 1)).astype(float))) x2 = [[1, 0], [2, 9], [0, 3]] try: @@ -54,7 +56,7 @@ def test_Array(): # input as numpy.ndarray x4 = np.array([[1, 0], [2, 9], [0, 3]]) - assert Array(x4, sparse=True).__str__() == ' (0, 0)\t1.0\n (1, 0)\t2.0\n (1, 1)\t9.0\n (2, 1)\t3.0' + assert_allclose_sparse(Array(x4, sparse=True), coo_array(x4.astype(float))) try: Array(x4, dim=1, sparse=True) except ValueError as e: @@ -62,7 +64,8 @@ def test_Array(): assert Array(np.array([1.0, 2.0, 3.0]), dtype=int, dim=2).__str__() == '[[1]\n [2]\n [3]]' assert Array(np.array([1.0, 2.0, 3.0]), dtype=int, dim=1).__str__() == '[1 2 3]' - assert Array(np.array([1.0, 2.0, 3.0]), - dtype=int, - dim=2, - sparse=True).__str__() == ' (0, 0)\t1\n (1, 0)\t2\n (2, 0)\t3' + assert_allclose_sparse(Array(np.array([1.0, 2.0, 3.0]), + dtype=int, + dim=2, + sparse=True), + coo_array(np.array([1.0, 2.0, 3.0]).reshape((-1, 1)).astype(float))) diff --git a/Solverz/solvers/daesolver/rodas/__init__.py b/Solverz/solvers/daesolver/rodas/__init__.py new file mode 100644 index 0000000..b4080d7 --- /dev/null +++ b/Solverz/solvers/daesolver/rodas/__init__.py @@ -0,0 +1 @@ +from Solverz.solvers.daesolver.rodas.rodas import Rodas diff --git a/Solverz/solvers/daesolver/rodas.py b/Solverz/solvers/daesolver/rodas/param.py similarity index 51% rename from Solverz/solvers/daesolver/rodas.py rename to Solverz/solvers/daesolver/rodas/param.py index ccc8394..dff6e45 100644 --- a/Solverz/solvers/daesolver/rodas.py +++ b/Solverz/solvers/daesolver/rodas/param.py @@ -1,306 +1,4 @@ -import warnings - -from Solverz.solvers.daesolver.utilities import * - - -@dae_io_parser -def Rodas(dae: nDAE, - tspan: List | np.ndarray, - y0: np.ndarray, - opt: Opt = None): - r""" - The stiffly accurate Rosenbrock methods including Rodas4 [1]_, Rodasp [2]_, Rodas5p [3]_. - - Parameters - ========== - - dae : nDAE - Numerical DAE object. - - tspan : List | np.ndarray - Either - - a list specifying [t0, tend], or - - a `np.ndarray` specifying the time nodes that you are concerned about - - y0 : np.ndarray - The initial values of variables - - opt : Opt - The solver options, including: - - - scheme: 'rodas4'(default)|'rodasp'|'rodas5p' - The rodas scheme - - rtol: 1e-3(default)|float - The relative error tolerance - - atol: 1e-6(default)|float - The absolute error tolerance - - event: Callable - The simulation events, with $t$ and $y$ being args, - and `value`, `is_terminal` and `direction` being outputs - - fixed_h: False(default)|bool - To use fixed step size - - hinit: None(default)|float - Initial step size - - hmax: None(default)|float - Maximum step size. - - Returns - ======= - - sol : daesol - The daesol object. - - - References - ========== - - .. [1] Hairer and Wanner, Solving Ordinary Differential Equations II , 2nd ed. Berlin, Heidelberg, Germany: Springer-Verlag, 1996. - .. [2] Steinebach, Order-reduction of ROW-methods for DAEs and method of lines applications. Preprint-Nr. 1741, FB Mathematik, TH Darmstadt, 1995 - .. [3] Steinebach, “Construction of rosenbrock-wanner method rodas5p and numerical benchmarks within the julia differential equations package,” BIT, vol. 63, no. 27, Jun 2023. - """ - - if opt is None: - opt = Opt() - stats = Stats(opt.scheme) - - rparam = Rodas_param(opt.scheme) - vsize = y0.shape[0] - tspan = np.array(tspan) - tend = tspan[-1] - t0 = tspan[0] - if opt.hmax is None: - opt.hmax = np.abs(tend - t0) - nt = 0 - t = t0 - hmin = 16 * np.spacing(t0) - uround = np.spacing(1.0) - T = np.zeros((10001,)) - T[nt] = t0 - y = np.zeros((10001, vsize)) - y0 = DaeIc(dae, y0, t0, opt.rtol) # check and modify initial values - y[0, :] = y0 - - dense_output = False - n_tspan = len(tspan) - told = t0 - if n_tspan > 2: - dense_output = True - inext = 1 - tnext = tspan[inext] - - events = opt.event - haveEvent = True if events is not None else False - - if haveEvent: - value, isterminal, direction = events(t, y0) - stop = 0 - nevent = -1 - te = np.zeros((10001,)) - ye = np.zeros((10001, vsize)) - ie = np.zeros((10001,)) - - # The initial step size - if opt.hinit is None: - dt = 1e-6 * (tend - t0) - else: - dt = opt.hinit - - dt = np.maximum(dt, hmin) - dt = np.minimum(dt, opt.hmax) - - M = dae.M - p = dae.p - done = False - reject = 0 - while not done: - # step size too small - # pass - - if np.abs(dt) < uround: - print(f"Error exit of RODAS at time = {t}: step size too small h = {dt}.\n") - stats.ret = 'failed' - break - - if reject > 100: - print(f"Step rejected over 100 times at time = {t}.\n") - stats.ret = 'failed' - break - - # Stretch the step if within 10% of T-t. - if t + dt >= tend: - dt = tend - t - else: - dt = np.minimum(dt, 0.5 * (tend - t)) - - if opt.fix_h: - dt = opt.hinit - - if done: - break - K = np.zeros((vsize, rparam.s)) - - if reject == 0: - J = dae.J(t, y0, p) - - dfdt0 = dt * dfdt(dae, t, y0) - rhs = dae.F(t, y0, p) + rparam.g[0] * dfdt0 - stats.nfeval = stats.nfeval + 1 - - try: - lu = lu_decomposition(M - dt * rparam.gamma * J) - except RuntimeError: - break - stats.ndecomp = stats.ndecomp + 1 - K[:, 0] = lu.solve(rhs) - - for j in range(1, rparam.s): - sum_1 = K @ rparam.alpha[:, j] - sum_2 = K @ rparam.gammatilde[:, j] - y1 = y0 + dt * sum_1 - - rhs = dae.F(t + dt * rparam.a[j], y1, p) + M @ sum_2 + rparam.g[j] * dfdt0 - stats.nfeval = stats.nfeval + 1 - sol = lu.solve(rhs) - K[:, j] = sol - sum_2 - - sum_1 = K @ (dt * rparam.b) - ynew = y0 + sum_1 - if not opt.fix_h: - sum_2 = K @ (dt * rparam.bd) - SK = (opt.atol + opt.rtol * abs(ynew)).reshape((-1,)) - err = np.max(np.abs((sum_1 - sum_2) / SK)) - if np.any(np.isinf(ynew)) or np.any(np.isnan(ynew)): - err = 1.0e6 - print('Warning Rodas: NaN or Inf occurs.') - err = np.maximum(err, 1.0e-6) - fac = opt.f_savety / (err ** (1 / rparam.pord)) - fac = np.minimum(opt.facmax, np.maximum(opt.fac1, fac)) - dtnew = dt * fac - else: - err = 1.0 - dtnew = dt - - if err <= 1.0: - reject = 0 - told = t - t = t + dt - stats.nstep = stats.nstep + 1 - # events - if haveEvent: - valueold = value - value, isterminal, direction = events(t, ynew) - value_save = value - ff = np.where(value * valueold < 0)[0] - if ff.size > 0: - for i in ff: - v0 = valueold[i] - v1 = value[i] - detect = 1 - if direction[i] < 0 and v0 <= v1: - detect = 0 - if direction[i] > 0 and v0 >= v1: - detect = 0 - if detect: - iterate = 1 - tL = told - tR = t - if np.abs(v1 - v0) > uround: - tevent = told - v0 * dt / (v1 - v0) # initial guess for tevent - else: - iterate = 0 - tevent = t - ynext = ynew - - tol = 128 * np.max([np.spacing(told), np.spacing(t)]) - tol = np.min([tol, np.abs(t - told)]) - while iterate > 0: - iterate = iterate + 1 - tau = (tevent - told) / dt - ynext = y0 + tau * dt * K @ ( - rparam.b + (tau - 1) * (rparam.c + tau * (rparam.d + tau * rparam.e))) - value, isterminal, direction = events(tevent, ynext) - if v1 * value[i] < 0: - tL = tevent - tevent = 0.5 * (tevent + tR) - v0 = value[i] - elif v0 * value[i] < 0: - tR = tevent - tevent = 0.5 * (tL + tevent) - v1 = value[i] - else: - iterate = 0 - if (tR - tL) < tol: - iterate = 0 - if iterate > 100: - print(f"Lost Event in interval [{told}, {t}].\n") - break - if np.abs(tevent - told) < opt.event_duration: - # We're not going to find events closer than tol. - break - t = tevent - ynew = ynext - nevent += 1 - te[nevent] = tevent - ie[nevent] = i - ye[nevent] = ynext - value, isterminal, direction = events(tevent, ynext) - value = value_save - if isterminal[i]: - if dense_output: - if tnext >= tevent: - tnext = tevent - stop = 1 - break - - if dense_output: # dense_output - while t >= tnext > told: - tau = (tnext - told) / dt - ynext = y0 + tau * dt * K @ (rparam.b + (tau - 1) * (rparam.c + tau * (rparam.d + tau * rparam.e))) - nt = nt + 1 - T[nt] = tnext - y[nt] = ynext - - if haveEvent and stop: - if tnext >= tevent: - break - - inext = inext + 1 - if inext <= n_tspan - 1: - tnext = tspan[inext] - if haveEvent and stop: - if tnext >= tevent: - tnext = tevent - else: - tnext = tend + dt - else: - nt = nt + 1 - T[nt] = t - y[nt] = ynew - - if nt == 10000: - warnings.warn("Time steps more than 10000! Rodas breaks. Try input a smaller tspan!") - done = True - - if np.abs(tend - t) < uround or stop: - done = True - y0 = ynew - opt.facmax = opt.fac2 - - else: - reject = reject + 1 - stats.nreject = stats.nreject + 1 - opt.facmax = 1 - dt = np.min([opt.hmax, np.max([hmin, dtnew])]) - - T = T[0:nt + 1] - y = y[0:nt + 1] - if haveEvent: - te = te[0:nevent + 1] - ye = ye[0:nevent + 1] - ie = ie[0:nevent + 1] - return daesol(T, y, te, ye, ie, stats) - else: - return daesol(T, y, stats=stats) +import numpy as np class Rodas_param: @@ -476,11 +174,3 @@ def __init__(self, self.gammatilde = self.gammatilde.T case _: raise ValueError("Not implemented") - - -def dfdt(dae: nDAE, t, y): - tscale = np.maximum(0.1 * np.abs(t), 1e-8) - ddt = t + np.sqrt(np.spacing(1)) * tscale - t - f0 = dae.F(t, y, dae.p) - f1 = dae.F(t + ddt, y, dae.p) - return (f1 - f0) / ddt diff --git a/Solverz/solvers/daesolver/rodas/rodas.py b/Solverz/solvers/daesolver/rodas/rodas.py new file mode 100644 index 0000000..22303ff --- /dev/null +++ b/Solverz/solvers/daesolver/rodas/rodas.py @@ -0,0 +1,312 @@ +import warnings + +from Solverz.solvers.daesolver.utilities import * +from Solverz.solvers.daesolver.rodas.param import Rodas_param + + +@dae_io_parser +def Rodas(dae: nDAE, + tspan: List | np.ndarray, + y0: np.ndarray, + opt: Opt = None): + r""" + The stiffly accurate Rosenbrock methods including Rodas4 [1]_, Rodasp [2]_, Rodas5p [3]_. + + Parameters + ========== + + dae : nDAE + Numerical DAE object. + + tspan : List | np.ndarray + Either + - a list specifying [t0, tend], or + - a `np.ndarray` specifying the time nodes that you are concerned about + + y0 : np.ndarray + The initial values of variables + + opt : Opt + The solver options, including: + + - scheme: 'rodas4'(default)|'rodasp'|'rodas5p' + The rodas scheme + - rtol: 1e-3(default)|float + The relative error tolerance + - atol: 1e-6(default)|float + The absolute error tolerance + - event: Callable + The simulation events, with $t$ and $y$ being args, + and `value`, `is_terminal` and `direction` being outputs + - fixed_h: False(default)|bool + To use fixed step size + - hinit: None(default)|float + Initial step size + - hmax: None(default)|float + Maximum step size. + + Returns + ======= + + sol : daesol + The daesol object. + + + References + ========== + + .. [1] Hairer and Wanner, Solving Ordinary Differential Equations II , 2nd ed. Berlin, Heidelberg, Germany: Springer-Verlag, 1996. + .. [2] Steinebach, Order-reduction of ROW-methods for DAEs and method of lines applications. Preprint-Nr. 1741, FB Mathematik, TH Darmstadt, 1995 + .. [3] Steinebach, “Construction of rosenbrock-wanner method rodas5p and numerical benchmarks within the julia differential equations package,” BIT, vol. 63, no. 27, Jun 2023. + """ + + if opt is None: + opt = Opt() + stats = Stats(opt.scheme) + + rparam = Rodas_param(opt.scheme) + vsize = y0.shape[0] + tspan = np.array(tspan) + tend = tspan[-1] + t0 = tspan[0] + if opt.hmax is None: + opt.hmax = np.abs(tend - t0) + nt = 0 + t = t0 + hmin = 16 * np.spacing(t0) + uround = np.spacing(1.0) + T = np.zeros((10001,)) + T[nt] = t0 + y = np.zeros((10001, vsize)) + y0 = DaeIc(dae, y0, t0, opt.rtol) # check and modify initial values + y[0, :] = y0 + + dense_output = False + n_tspan = len(tspan) + told = t0 + if n_tspan > 2: + dense_output = True + inext = 1 + tnext = tspan[inext] + + events = opt.event + haveEvent = True if events is not None else False + + if haveEvent: + value, isterminal, direction = events(t, y0) + stop = 0 + nevent = -1 + te = np.zeros((10001,)) + ye = np.zeros((10001, vsize)) + ie = np.zeros((10001,)) + + # The initial step size + if opt.hinit is None: + dt = 1e-6 * (tend - t0) + else: + dt = opt.hinit + + dt = np.maximum(dt, hmin) + dt = np.minimum(dt, opt.hmax) + + M = dae.M + p = dae.p + done = False + reject = 0 + while not done: + # step size too small + # pass + + if np.abs(dt) < uround: + print(f"Error exit of RODAS at time = {t}: step size too small h = {dt}.\n") + stats.ret = 'failed' + break + + if reject > 100: + print(f"Step rejected over 100 times at time = {t}.\n") + stats.ret = 'failed' + break + + # Stretch the step if within 10% of T-t. + if t + dt >= tend: + dt = tend - t + else: + dt = np.minimum(dt, 0.5 * (tend - t)) + + if opt.fix_h: + dt = opt.hinit + + if done: + break + K = np.zeros((vsize, rparam.s)) + + if reject == 0: + J = dae.J(t, y0, p) + + dfdt0 = dt * dfdt(dae, t, y0) + rhs = dae.F(t, y0, p) + rparam.g[0] * dfdt0 + stats.nfeval = stats.nfeval + 1 + + try: + lu = lu_decomposition(M - dt * rparam.gamma * J) + except RuntimeError: + break + stats.ndecomp = stats.ndecomp + 1 + K[:, 0] = lu.solve(rhs) + + for j in range(1, rparam.s): + sum_1 = K @ rparam.alpha[:, j] + sum_2 = K @ rparam.gammatilde[:, j] + y1 = y0 + dt * sum_1 + + rhs = dae.F(t + dt * rparam.a[j], y1, p) + M @ sum_2 + rparam.g[j] * dfdt0 + stats.nfeval = stats.nfeval + 1 + sol = lu.solve(rhs) + K[:, j] = sol - sum_2 + + sum_1 = K @ (dt * rparam.b) + ynew = y0 + sum_1 + if not opt.fix_h: + sum_2 = K @ (dt * rparam.bd) + SK = (opt.atol + opt.rtol * abs(ynew)).reshape((-1,)) + err = np.max(np.abs((sum_1 - sum_2) / SK)) + if np.any(np.isinf(ynew)) or np.any(np.isnan(ynew)): + err = 1.0e6 + print('Warning Rodas: NaN or Inf occurs.') + err = np.maximum(err, 1.0e-6) + fac = opt.f_savety / (err ** (1 / rparam.pord)) + fac = np.minimum(opt.facmax, np.maximum(opt.fac1, fac)) + dtnew = dt * fac + else: + err = 1.0 + dtnew = dt + + if err <= 1.0: + reject = 0 + told = t + t = t + dt + stats.nstep = stats.nstep + 1 + # events + if haveEvent: + valueold = value + value, isterminal, direction = events(t, ynew) + value_save = value + ff = np.where(value * valueold < 0)[0] + if ff.size > 0: + for i in ff: + v0 = valueold[i] + v1 = value[i] + detect = 1 + if direction[i] < 0 and v0 <= v1: + detect = 0 + if direction[i] > 0 and v0 >= v1: + detect = 0 + if detect: + iterate = 1 + tL = told + tR = t + if np.abs(v1 - v0) > uround: + tevent = told - v0 * dt / (v1 - v0) # initial guess for tevent + else: + iterate = 0 + tevent = t + ynext = ynew + + tol = 128 * np.max([np.spacing(told), np.spacing(t)]) + tol = np.min([tol, np.abs(t - told)]) + while iterate > 0: + iterate = iterate + 1 + tau = (tevent - told) / dt + ynext = y0 + tau * dt * K @ ( + rparam.b + (tau - 1) * (rparam.c + tau * (rparam.d + tau * rparam.e))) + value, isterminal, direction = events(tevent, ynext) + if v1 * value[i] < 0: + tL = tevent + tevent = 0.5 * (tevent + tR) + v0 = value[i] + elif v0 * value[i] < 0: + tR = tevent + tevent = 0.5 * (tL + tevent) + v1 = value[i] + else: + iterate = 0 + if (tR - tL) < tol: + iterate = 0 + if iterate > 100: + print(f"Lost Event in interval [{told}, {t}].\n") + break + if np.abs(tevent - told) < opt.event_duration: + # We're not going to find events closer than tol. + break + t = tevent + ynew = ynext + nevent += 1 + te[nevent] = tevent + ie[nevent] = i + ye[nevent] = ynext + value, isterminal, direction = events(tevent, ynext) + value = value_save + if isterminal[i]: + if dense_output: + if tnext >= tevent: + tnext = tevent + stop = 1 + break + + if dense_output: # dense_output + while t >= tnext > told: + tau = (tnext - told) / dt + ynext = y0 + tau * dt * K @ (rparam.b + (tau - 1) * (rparam.c + tau * (rparam.d + tau * rparam.e))) + nt = nt + 1 + T[nt] = tnext + y[nt] = ynext + + if haveEvent and stop: + if tnext >= tevent: + break + + inext = inext + 1 + if inext <= n_tspan - 1: + tnext = tspan[inext] + if haveEvent and stop: + if tnext >= tevent: + tnext = tevent + else: + tnext = tend + dt + else: + nt = nt + 1 + T[nt] = t + y[nt] = ynew + + if nt == 10000: + warnings.warn("Time steps more than 10000! Rodas breaks. Try input a smaller tspan!") + done = True + + if np.abs(tend - t) < uround or stop: + done = True + y0 = ynew + opt.facmax = opt.fac2 + + else: + reject = reject + 1 + stats.nreject = stats.nreject + 1 + opt.facmax = 1 + dt = np.min([opt.hmax, np.max([hmin, dtnew])]) + + T = T[0:nt + 1] + y = y[0:nt + 1] + if haveEvent: + te = te[0:nevent + 1] + ye = ye[0:nevent + 1] + ie = ie[0:nevent + 1] + return daesol(T, y, te, ye, ie, stats) + else: + return daesol(T, y, stats=stats) + + +def dfdt(dae: nDAE, t, y): + tscale = np.maximum(0.1 * np.abs(t), 1e-8) + ddt = t + np.sqrt(np.spacing(1)) * tscale - t + f0 = dae.F(t, y, dae.p) + f1 = dae.F(t + ddt, y, dae.p) + return (f1 - f0) / ddt diff --git a/Solverz/solvers/daesolver/trapezoidal.py b/Solverz/solvers/daesolver/trapezoidal.py index 9dc9ae6..5333058 100644 --- a/Solverz/solvers/daesolver/trapezoidal.py +++ b/Solverz/solvers/daesolver/trapezoidal.py @@ -66,9 +66,8 @@ def implicit_trapezoid(dae: nDAE, sol = nr_method(ae, y0, Opt(stats=True)) y1 = sol.y - ite = sol.stats - stats.ndecomp = stats.ndecomp + ite - stats.nfeval = stats.nfeval + ite + stats.ndecomp = stats.ndecomp + sol.stats.nstep + stats.nfeval = stats.nfeval + sol.stats.nstep tt = tt + dt nt = nt + 1 diff --git a/Solverz/solvers/fdesolver.py b/Solverz/solvers/fdesolver.py index 68d5719..0a065c4 100644 --- a/Solverz/solvers/fdesolver.py +++ b/Solverz/solvers/fdesolver.py @@ -90,10 +90,9 @@ def fdae_solver(fdae: nFDAE, sol = nr_method(ae, u0, Opt(ite_tol=opt.ite_tol, stats=True)) u1 = sol.y - ite = sol.stats - stats.ndecomp = stats.ndecomp + ite - stats.nfeval = stats.nfeval + ite + 1 - if ite >= 100: + stats.ndecomp = stats.ndecomp + sol.stats.ndecomp + stats.nfeval = stats.nfeval + stats.nfeval + if stats.nstep >= 100: print(f"FDAE solver broke at time={tt} due to non-convergence") break diff --git a/Solverz/solvers/nlaesolver/__init__.py b/Solverz/solvers/nlaesolver/__init__.py index 5329621..babf93c 100644 --- a/Solverz/solvers/nlaesolver/__init__.py +++ b/Solverz/solvers/nlaesolver/__init__.py @@ -1,3 +1,4 @@ from Solverz.solvers.nlaesolver.nr import nr_method from Solverz.solvers.nlaesolver.cnr import continuous_nr - +from Solverz.solvers.nlaesolver.sicnm import sicnm +from Solverz.solvers.nlaesolver.lm import lm diff --git a/Solverz/solvers/nlaesolver/cnr.py b/Solverz/solvers/nlaesolver/cnr.py index e9dedca..04bcb44 100644 --- a/Solverz/solvers/nlaesolver/cnr.py +++ b/Solverz/solvers/nlaesolver/cnr.py @@ -20,7 +20,7 @@ def continuous_nr(eqn: nAE, opt : Opt The solver options, including: - - ite_tol: 1e-8(default)|float + - ite_tol: 1e-5(default)|float The iteration error tolerance. Returns @@ -40,6 +40,8 @@ def continuous_nr(eqn: nAE, if opt is None: opt = Opt() + stats = Stats('CNM based on ERK4') + dt = 1 hmax = 1.7 tol = opt.ite_tol @@ -66,6 +68,7 @@ def f(y_, p_) -> np.ndarray: # atol and rtol can not be too small ite = 0 df = eqn.F(y, p) + stats.nfeval += 1 while max(abs(df)) > tol: ite = ite + 1 err = 2 @@ -81,6 +84,9 @@ def f(y_, p_) -> np.ndarray: k7 = f(ynew, p) kE = k1 * e1 + k3 * e3 + k4 * e4 + k5 * e5 + k6 * e6 + k7 * e7 + stats.nfeval += 7 + stats.ndecomp += 7 + # error control # error estimation err = dt * np.linalg.norm( @@ -99,6 +105,7 @@ def f(y_, p_) -> np.ndarray: y = ynew df = eqn.F(y, p) + stats.nfeval += 1 if nofailed: # Enlarge step size if no failure is met temp = 1.25 * (err / rtol) ** Pow @@ -110,6 +117,12 @@ def f(y_, p_) -> np.ndarray: dt = np.min([dt, hmax]) if ite > 100: + print(f"Cannot converge within 100 iterations. Deviation: {max(abs(df))}!") + stats.succeed = False break - return aesol(y, ite) + if np.any(np.isnan(y)): + stats.succeed = False + stats.nstep = ite + + return aesol(y, stats) diff --git a/Solverz/solvers/nlaesolver/lm.py b/Solverz/solvers/nlaesolver/lm.py index d6fe86e..11bb577 100644 --- a/Solverz/solvers/nlaesolver/lm.py +++ b/Solverz/solvers/nlaesolver/lm.py @@ -23,7 +23,7 @@ def lm(eqn: nAE, opt : Opt The solver options, including: - - ite_tol: 1e-8(default)|float + - ite_tol: 1e-5(default)|float The iteration error tolerance. Returns @@ -40,12 +40,16 @@ def lm(eqn: nAE, """ if opt is None: - opt = Opt(ite_tol=1e-8) + opt = Opt() + + stats = Stats('Levenberg–Marquardt') tol = opt.ite_tol p = eqn.p # optimize.root func cannot handle callable jac that returns scipy.sparse.csc_array sol = optimize.root(lambda x: eqn.F(x, p), y, jac=lambda x: eqn.J(x, p).toarray(), method='lm', tol=tol) + stats.succeed = sol.success + stats.nfeval = sol.nfev - return aesol(sol.x, sol.njev) + return aesol(sol.x, stats) diff --git a/Solverz/solvers/nlaesolver/nr.py b/Solverz/solvers/nlaesolver/nr.py index eed8408..ad47666 100644 --- a/Solverz/solvers/nlaesolver/nr.py +++ b/Solverz/solvers/nlaesolver/nr.py @@ -20,7 +20,7 @@ def nr_method(eqn: nAE, opt : Opt The solver options, including: - - ite_tol: 1e-8(default)|float + - ite_tol: 1e-5(default)|float The iteration error tolerance. Returns @@ -36,19 +36,29 @@ def nr_method(eqn: nAE, """ if opt is None: - opt = Opt(ite_tol=1e-8) + opt = Opt() tol = opt.ite_tol p = eqn.p df = eqn.F(y, p) - ite = 0 + + stats = Stats('Newton') + stats.nfeval += 1 + # main loop while max(abs(df)) > tol: - ite = ite + 1 y = y - solve(eqn.J(y, p), df) + stats.ndecomp += 1 df = eqn.F(y, p) - if ite >= 100: + stats.nfeval += 1 + stats.nstep += 1 + + if stats.nstep >= 100: print(f"Cannot converge within 100 iterations. Deviation: {max(abs(df))}!") + stats.succeed = False break - return aesol(y, ite) + if np.any(np.isnan(y)): + stats.succeed = False + + return aesol(y, stats) diff --git a/Solverz/solvers/nlaesolver/sicnm.py b/Solverz/solvers/nlaesolver/sicnm.py new file mode 100644 index 0000000..9455e37 --- /dev/null +++ b/Solverz/solvers/nlaesolver/sicnm.py @@ -0,0 +1,369 @@ +import warnings + +from Solverz.solvers.nlaesolver.utilities import * +from Solverz.solvers.daesolver.rodas.param import Rodas_param +from scipy.sparse import eye_array as speye +from scipy.sparse.linalg import splu, spsolve_triangular +from scipy.sparse import csc_array, block_array +from Solverz.solvers.parser import dae_io_parser +from Solverz.solvers.solution import daesol + + +@ae_io_parser +def sicnm(ae: nAE, + y0: np.ndarray, + opt: Opt = None): + r""" + The semi-implicit continuous Newton method [1]_. The philosophy is to rewrite the algebraic equation + + + .. math:: + + 0=g(y) + + as the differential algebraic equations + + .. math:: + + \begin{aligned} + \dot{y}&=z \\ + 0&=J(y)z+g(y) + \end{aligned}, + + with $y_0$ being the initial value guess and $z_0=-J(y_0)^{-1}g(y_0)$, where $z$ is an intermediate variable introduced. Then the DAEs are solved by Rodas. SICNM is found to be more robust than the Newton's method, for which the theoretical proof can be found in my paper [1]_. In addition, the non-iterative nature of Rodas guarantees the efficiency. + + One can change the rodas scheme according to the ones implemented in the DAE version of Rodas. + + SICNM used the Hessian-vector product (HVP) interface of the algebraic equation models, the `make_hvp` flag should be set to `True` when printing numerical modules. + + An illustrative example of SICNM can be found in the `power flow section `_ of Solverz' cookbook. + + Parameters + ========== + + eqn : nAE + Numerical AE object. + + y : np.ndarray + The initial values of variables + + opt : Opt + The solver options, including: + + - ite_tol: 1e-8(default)|float + The iteration error tolerance. + + Returns + ======= + + sol : aesol + The aesol object. + + References + ========== + + .. [1] R. Yu, W. Gu, S. Lu, and Y. Xu, “Semi-implicit continuous newton method for power flow analysis,” 2023, arXiv:2312.02809. + + + """ + p = ae.p + + def F_tilda(y_, v_): + return np.concatenate([v_, ae.J(y_, p) @ v_ + ae.F(y_, p)]) + + if opt is None: + opt = Opt() + stats = Stats('Sicnm based on ' + opt.scheme) + + rparam = Rodas_param(opt.scheme) + vsize = y0.shape[0] + tspan = np.array([0, 10000]) + tend = tspan[-1] + t0 = tspan[0] + if opt.hmax is None: + opt.hmax = np.abs(tend - t0) + nt = 0 + t = t0 + hmin = 16 * np.spacing(t0) + uround = np.spacing(1.0) + T = np.zeros((10001,)) + T[nt] = t0 + y = np.zeros((10001, vsize)) + v = np.zeros((10001, vsize)) + y[0, :] = y0 + v0 = solve(ae.J(y0, p), -ae.F(y0, p)) + v[0, :] = v0 + + dense_output = False + n_tspan = len(tspan) + told = t0 + if n_tspan > 2: + dense_output = True + inext = 1 + tnext = tspan[inext] + + events = opt.event + haveEvent = True if events is not None else False + + if haveEvent: + value, isterminal, direction = events(t, y0) + stop = 0 + nevent = -1 + te = np.zeros((10001,)) + ye = np.zeros((10001, vsize)) + ie = np.zeros((10001,)) + + # The initial step size + if opt.hinit is None: + dt = 1e-6 * (tend - t0) + else: + dt = opt.hinit + + dt = np.maximum(dt, hmin) + dt = np.minimum(dt, opt.hmax) + + ZERO = None + EYE = speye(vsize, format='csc') + M = csc_array((np.ones(vsize), (np.arange(0, vsize), np.arange(0, vsize))), + shape=(2 * vsize, 2 * vsize)) + done = False + reject = 0 + while not done: + # step size too small + # pass + + if np.abs(dt) < uround: + print(f"Error exit of RODAS at time = {t}: step size too small h = {dt}.\n") + stats.ret = 'failed' + stats.succeed = False + break + + if reject > 100: + print(f"Step rejected over 100 times at time = {t}.\n") + stats.ret = 'failed' + stats.succeed = False + break + + # Stretch the step if within 10% of T-t. + if t + dt >= tend: + dt = tend - t + else: + dt = np.minimum(dt, 0.5 * (tend - t)) + + if done: + break + K = np.zeros((2 * vsize, rparam.s)) + + if reject == 0: + Hvp = ae.HVP(y0, p, v0) + J = ae.J(y0, p) + + dfdt0 = 0 + rhs = F_tilda(y0, v0) + rparam.g[0] * dfdt0 + stats.nfeval = stats.nfeval + 1 + + try: + if opt.partial_decompose: + # partial decomposition + N = - dt * rparam.gamma * (Hvp + J) + Lambda = dt * rparam.gamma * (N - J) + lu = splu(Lambda) + if nt == 0: + P = csc_array((np.ones(vsize), (lu.perm_r, np.arange(vsize)))) + Q = csc_array((np.ones(vsize), (np.arange(vsize), lu.perm_c))) + # b_perm = np.concatenate([np.arange(vsize), lu.perm_r + vsize]) + # dx_perm = np.concatenate([np.arange(vsize), lu.perm_c + vsize]) + P_tilda = block_array([[EYE, ZERO], [ZERO, P]], format='csc') + Q_tilda = block_array([[EYE, ZERO], [ZERO, Q]], format='csc') + + L_tilda = block_array([[EYE, ZERO], [P @ N, lu.L]], format='csc') + U_tilda = block_array([[EYE, -dt * rparam.gamma * Q], [ZERO, lu.U]], format='csc') + else: + # full decomposition + tilde_J = block_array([[ZERO, EYE], [Hvp + J, J]]) + tilde_E = M - dt * rparam.gamma * tilde_J + lu = splu(tilde_E) + except RuntimeError: + stats.succeed = False + break + + stats.ndecomp = stats.ndecomp + 1 + if opt.partial_decompose: + # partial decomposition + # K[dx_perm, 0] = solve(U_tilda, solve(L_tilda, rhs[b_perm])) + K[:, 0] = Q_tilda@(solve(U_tilda, solve(L_tilda, P_tilda@rhs))) + # not stable be very careful with the following triangular solver + # K[:, 0] = Q_tilda@(spsolve_triangular(U_tilda, spsolve_triangular(L_tilda, P_tilda@rhs), False)) + else: + # full decomposition + K[:, 0] = lu.solve(rhs) + # K[:, 0] = solve(tilde_E, rhs) + + for j in range(1, rparam.s): + sum_1 = K @ rparam.alpha[:, j] + sum_2 = K @ rparam.gammatilde[:, j] + y1 = y0 + dt * sum_1[0:vsize] + v1 = v0 + dt * sum_1[vsize:2 * vsize] + + rhs = F_tilda(y1, v1) + M @ sum_2 + rparam.g[j] * dfdt0 + stats.nfeval = stats.nfeval + 1 + + if opt.partial_decompose: + # partial decomposition + # K[dx_perm, j] = solve(U_tilda, solve(L_tilda, rhs[b_perm])) + K[:, j] = Q_tilda @ (solve(U_tilda, solve(L_tilda, P_tilda @ rhs))) + # not stable be very careful with the following triangular solver + # K[:, j] = Q_tilda @ (spsolve_triangular(U_tilda, spsolve_triangular(L_tilda, P_tilda @ rhs), False)) + else: + # full decomposition + K[:, j] = lu.solve(rhs) + # K[:, j] = solve(tilde_E, rhs) + + K[:, j] = K[:, j] - sum_2 + + sum_1 = K @ (dt * rparam.b) + ynew = y0 + sum_1[0:vsize] + vnew = v0 + sum_1[vsize:2 * vsize] + sum_2 = K @ (dt * rparam.bd) + y_tilda = np.concatenate([ynew, vnew]) + SK = (opt.atol + opt.rtol * abs(y_tilda)).reshape((-1,)) + err = np.max(np.abs((sum_1 - sum_2) / SK)) + if np.any(np.isinf(y_tilda)) or np.any(np.isnan(y_tilda)): + err = 1.0e6 + print('Warning Rodas: NaN or Inf occurs.') + err = np.maximum(err, 1.0e-6) + fac = opt.f_savety / (err ** (1 / rparam.pord)) + fac = np.minimum(opt.facmax, np.maximum(opt.fac1, fac)) + dtnew = dt * fac + + if err <= 1.0: + reject = 0 + t = t + dt + stats.nstep = stats.nstep + 1 + # events + if haveEvent: + valueold = value + value, isterminal, direction = events(t, ynew) + value_save = value + ff = np.where(value * valueold < 0)[0] + if ff.size > 0: + for i in ff: + v0_ = valueold[i] + v1_ = value[i] + detect = 1 + if direction[i] < 0 and v0_ <= v1_: + detect = 0 + if direction[i] > 0 and v0_ >= v1_: + detect = 0 + if detect: + iterate = 1 + tL = told + tR = t + if np.abs(v1_ - v0_) > uround: + tevent = told - v0_ * dt / (v1_ - v0_) # initial guess for tevent + else: + iterate = 0 + tevent = t + ynext = ynew + + tol = 128 * np.max([np.spacing(told), np.spacing(t)]) + tol = np.min([tol, np.abs(t - told)]) + while iterate > 0: + iterate = iterate + 1 + tau = (tevent - told) / dt + ynext = y0 + (tau * dt * K @ (rparam.b + (tau - 1) * (rparam.c + tau * (rparam.d + tau * rparam.e))))[0:vsize] + value, isterminal, direction = events(tevent, ynext) + if v1_ * value[i] < 0: + tL = tevent + tevent = 0.5 * (tevent + tR) + v0_ = value[i] + elif v0_ * value[i] < 0: + tR = tevent + tevent = 0.5 * (tL + tevent) + v1_ = value[i] + else: + iterate = 0 + if (tR - tL) < tol: + iterate = 0 + if iterate > 100: + print(f"Lost Event in interval [{told}, {t}].\n") + break + if np.abs(tevent - told) < opt.event_duration: + # We're not going to find events closer than tol. + break + t = tevent + ynew = ynext + nevent += 1 + te[nevent] = tevent + ie[nevent] = i + ye[nevent] = ynext + value, isterminal, direction = events(tevent, ynext) + value = value_save + if isterminal[i]: + if dense_output: + if tnext >= tevent: + tnext = tevent + stop = 1 + break + + if dense_output: # dense_output + while t >= tnext > told: + tau = (tnext - told) / dt + ynext = y0 + tau * dt * K @ (rparam.b + (tau - 1) * (rparam.c + tau * (rparam.d + tau * rparam.e))) + nt = nt + 1 + T[nt] = tnext + y[nt] = ynext + + if haveEvent and stop: + if tnext >= tevent: + break + + inext = inext + 1 + if inext <= n_tspan - 1: + tnext = tspan[inext] + if haveEvent and stop: + if tnext >= tevent: + tnext = tevent + else: + tnext = tend + dt + else: + nt = nt + 1 + T[nt] = t + y[nt] = ynew + v[nt] = vnew + + if nt == 10000: + warnings.warn("Time steps more than 10000! Rodas breaks. Try input a smaller tspan!") + stats.succeed = False + done = True + + if np.abs(tend - t) < uround: + done = True + if np.max(np.abs(ae.F(ynew, p))) < opt.ite_tol: + done = True + if stop: + done = True + + y0 = ynew + v0 = vnew + opt.facmax = opt.fac2 + + else: + reject = reject + 1 + stats.nreject = stats.nreject + 1 + opt.facmax = 1 + dt = np.min([opt.hmax, np.max([hmin, dtnew])]) + + T = T[0:nt + 1] + y = y[0:nt + 1] + + sol = aesol(y[-1], stats=stats) + + stats.T = T + stats.y = y + if haveEvent: + stats.te = te[0:nevent + 1] + stats.ye = ye[0:nevent + 1] + stats.ie = ie[0:nevent + 1] + + return sol diff --git a/Solverz/solvers/nlaesolver/utilities.py b/Solverz/solvers/nlaesolver/utilities.py index d9fb9c1..f0432de 100644 --- a/Solverz/solvers/nlaesolver/utilities.py +++ b/Solverz/solvers/nlaesolver/utilities.py @@ -8,3 +8,4 @@ from Solverz.solvers.option import Opt from Solverz.solvers.parser import ae_io_parser from Solverz.solvers.solution import aesol +from Solverz.solvers.stats import Stats diff --git a/Solverz/solvers/option.py b/Solverz/solvers/option.py index 4ae667e..e6685f2 100644 --- a/Solverz/solvers/option.py +++ b/Solverz/solvers/option.py @@ -18,7 +18,8 @@ def __init__(self, stats=False, step_size=1e-3, event: Callable = None, - event_duration=1e-8): + event_duration=1e-8, + partial_decompose=False): self.atol = atol self.rtol = rtol self.f_savety = f_savety @@ -35,3 +36,4 @@ def __init__(self, self.step_size = step_size self.event = event self.event_duration = event_duration + self.partial_decompose = partial_decompose diff --git a/Solverz/solvers/solution.py b/Solverz/solvers/solution.py index c548acf..90330a1 100644 --- a/Solverz/solvers/solution.py +++ b/Solverz/solvers/solution.py @@ -15,6 +15,9 @@ def __init__(self, self.y = y self.stats = stats + def __repr__(self): + return f'ae solution using {self.stats.scheme}, succeed: {self.stats.succeed}' + class daesol: __slots__ = ['T', 'Y', 'te', 'ie', 'ye', 'stats'] diff --git a/Solverz/solvers/stats.py b/Solverz/solvers/stats.py index 24f6f2d..3f160d3 100644 --- a/Solverz/solvers/stats.py +++ b/Solverz/solvers/stats.py @@ -7,6 +7,12 @@ def __init__(self, scheme=None): self.ndecomp = 0 self.nreject = 0 self.ret = None + self.succeed = True def __repr__(self): - return f"Scheme {self.scheme}, nstep: {self.nstep}, nfeval: {self.nfeval}, ndecomp: {self.ndecomp}, nreject: {self.nreject}." + return (f"Scheme: {self.scheme}, " + f"succeed: {self.succeed}, " + f"nstep: {self.nstep}, " + f"nfeval: {self.nfeval}, " + f"ndecomp: {self.ndecomp}, " + f"nreject: {self.nreject}.") diff --git a/Solverz/sym_algebra/functions.py b/Solverz/sym_algebra/functions.py index 23492c3..3133b6e 100644 --- a/Solverz/sym_algebra/functions.py +++ b/Solverz/sym_algebra/functions.py @@ -114,6 +114,16 @@ def _numpycode(self, printer, **kwargs): temp += '@({operand})'.format(operand=printer._print(arg)) return r'(' + temp + r')' + def _octave(self, printer, **kwargs): + + temp = printer._print(self.args[0]) + for arg in self.args[1:]: + if isinstance(arg, (Symbol, Function)): + temp += '*{operand}'.format(operand=printer._print(arg)) + else: + temp += '*({operand})'.format(operand=printer._print(arg)) + return r'(' + temp + r')' + def _lambdacode(self, printer, **kwargs): return self._numpycode(printer, **kwargs) @@ -157,6 +167,9 @@ def _lambdacode(self, printer, **kwargs): def _pythoncode(self, printer, **kwargs): return self._numpycode(printer, **kwargs) + def _octave(self, printer, **kwargs): + return r'diag(' + printer._print(self.args[0], **kwargs) + r')' + # %% Univariate func @VarParser @@ -214,6 +227,21 @@ def _pythoncode(self, printer, **kwargs): return self._numpycode(printer, **kwargs) +class ln(UniVarFunc): + r""" + The ln function, $ln(x)$. + """ + + def fdiff(self, argindex=1): + return 1 / self.args[0] + + def _numpycode(self, printer, **kwargs): + return r'log(' + printer._print(self.args[0]) + r')' + + def _pythoncode(self, printer, **kwargs): + return self._numpycode(printer, **kwargs) + + class sin(UniVarFunc): r""" The sine function. @@ -280,6 +308,43 @@ def _lambdacode(self, printer, **kwargs): def _pythoncode(self, printer, **kwargs): return self._numpycode(printer, **kwargs) + def _octave(self, printer, **kwargs): + return r'sign(' + printer._print(self.args[0], **kwargs) + r')' + + +class heaviside(UniVarFunc): + r""" + The heaviside step function + + .. math:: + + \operatorname{Heaviside}(x)= + \begin{cases} + 1 & x >= 0\\ + 0 & x < 0\\ + \end{cases} + + which should be distinguished from sympy.Heaviside + """ + + def fdiff(self, argindex=1): + # sign function should be treated as a constant. + if argindex == 1: + return 0 + raise ArgumentIndexError(self, argindex) + + def _numpycode(self, printer, **kwargs): + return r'Heaviside(' + printer._print(self.args[0], **kwargs) + r')' + + def _lambdacode(self, printer, **kwargs): + return self._numpycode(printer, **kwargs) + + def _pythoncode(self, printer, **kwargs): + return self._numpycode(printer, **kwargs) + + def _octave(self, printer, **kwargs): + return r'Heaviside(' + printer._print(self.args[0], **kwargs) + r')' + # %% multi-variate func @VarParser diff --git a/Solverz/sym_algebra/matrix_calculus.py b/Solverz/sym_algebra/matrix_calculus.py index 11c12cb..1329a06 100644 --- a/Solverz/sym_algebra/matrix_calculus.py +++ b/Solverz/sym_algebra/matrix_calculus.py @@ -2,7 +2,6 @@ from typing import List, Union, Dict, Tuple -import matplotlib.pyplot as plt import networkx as nx import numpy as np import sympy as sp @@ -331,6 +330,7 @@ def draw_dag_as_tree(G, pos=None): nx.draw_networkx_labels(G, pos) nx.draw_networkx_edges(G, pos, arrowstyle='-|>', arrowsize=15) + import matplotlib.pyplot as plt plt.axis('off') plt.show() diff --git a/Solverz/sym_algebra/symbols.py b/Solverz/sym_algebra/symbols.py index edd8ab8..37bf5cd 100644 --- a/Solverz/sym_algebra/symbols.py +++ b/Solverz/sym_algebra/symbols.py @@ -66,7 +66,9 @@ def SymbolExtractor(index) -> Dict: return temp -Solverz_internal_name = ['y_', 'F_', 'p_', 'J_', 'row', 'col', 'data', '_F_', 'data_'] +Solverz_internal_name = ['y_', 'F_', 'p_', 'J_', + 'row', 'col', 'data', '_F_', 'data_', 'Hvp_', 'v_', 'row_hvp', 'col_hvp', 'data_hvp', + '_data_', '_data_hvp'] class SolSymBasic(Symbol): @@ -78,7 +80,8 @@ class SolSymBasic(Symbol): def __new__(cls, name: str, value=None, dim: int = 1, internal_use=False): if any([name == built_in_name for built_in_name in Solverz_internal_name]): if not internal_use: - raise ValueError(f"Solverz built-in name {name}, cannot be used as variable name.") + raise ValueError( + f"Solverz built-in name {name}, cannot be used as variable name.") obj = Symbol.__new__(cls, f'{name}') obj.name = f'{name}' obj.dim = dim @@ -127,13 +130,18 @@ def IndexCodePrinter(index, printer): stop = index.stop step = index.step if any([isinstance(arg, (idx, Expr)) for arg in [start, stop, step]]): - start = IndexCodePrinter(start, printer) if start is not None else None - stop = IndexCodePrinter(stop, printer) if stop is not None else None - step = IndexCodePrinter(step, printer) if step is not None else None + start = IndexCodePrinter( + start, printer) if start is not None else None + stop = IndexCodePrinter( + stop, printer) if stop is not None else None + step = IndexCodePrinter( + step, printer) if step is not None else None return 'sol_slice({i}, {j}, {k})'.format(i=start, j=stop, k=step) else: - start = IndexCodePrinter(start, printer) if start is not None else '' - stop = IndexCodePrinter(stop, printer) if stop is not None else '' + start = IndexCodePrinter( + start, printer) if start is not None else '' + stop = IndexCodePrinter( + stop, printer) if stop is not None else '' slice_str = '{i}:{j}'.format(i=start, j=stop) slice_str += f':{IndexCodePrinter(step, printer)}' if step is not None else '' return slice_str @@ -142,7 +150,8 @@ def IndexCodePrinter(index, printer): if len(self.index) != 2: raise ValueError("Support only two element tuples!") else: - temp = IndexCodePrinter(self.index[0], printer) + ',' + IndexCodePrinter(self.index[1], printer) + temp = IndexCodePrinter( + self.index[0], printer) + ',' + IndexCodePrinter(self.index[1], printer) else: temp = IndexCodePrinter(self.index, printer) diff --git a/Solverz/utilities/test/test_typechecker.py b/Solverz/utilities/test/test_typechecker.py new file mode 100644 index 0000000..6487d07 --- /dev/null +++ b/Solverz/utilities/test/test_typechecker.py @@ -0,0 +1,7 @@ +import numpy as np +from Solverz.utilities.type_checker import is_integer + + +def test_is_number(): + assert is_integer(np.array([1.0]).astype(int)[0]) + assert not is_integer(np.array([1.0])[0]) diff --git a/Solverz/utilities/testing.py b/Solverz/utilities/testing.py new file mode 100644 index 0000000..a8ff048 --- /dev/null +++ b/Solverz/utilities/testing.py @@ -0,0 +1,14 @@ +import numpy as np +from scipy.sparse import coo_array, sparray + + +def assert_allclose_sparse(a, b, rtol=1e-7, atol=0): + if not isinstance(a, sparray): + raise TypeError(f"a should be scipy.sparray, instead of {type(a)}") + if not isinstance(b, sparray): + raise TypeError(f"b should be scipy.sparray, instead of {type(b)}") + a_ = coo_array(a) + b_ = coo_array(b) + np.testing.assert_allclose(a_.data, b_.data, rtol, atol) + np.testing.assert_allclose(a_.row, b_.row, rtol, atol) + np.testing.assert_allclose(a_.col, b_.col, rtol, atol) diff --git a/Solverz/utilities/type_checker.py b/Solverz/utilities/type_checker.py index d459962..74ecfcb 100644 --- a/Solverz/utilities/type_checker.py +++ b/Solverz/utilities/type_checker.py @@ -1,5 +1,6 @@ from sympy import Number as SymNumber, Integer import numpy as np +from numpy import integer as NpInteger from numbers import Number as PyNumber @@ -13,7 +14,7 @@ def is_number(num): def is_integer(num): - if isinstance(num, (int, Integer)): + if isinstance(num, (int, Integer, NpInteger)): return True else: return False @@ -32,3 +33,7 @@ def is_scalar(a): return True else: return False + + +def is_zero(a): + return a == 0 diff --git a/docs/src/gettingstart.md b/docs/src/gettingstart.md index 800fccd..6cbaa49 100644 --- a/docs/src/gettingstart.md +++ b/docs/src/gettingstart.md @@ -203,7 +203,25 @@ def J_(y_, p_): return coo_array((data, (row,col)), (1, 1)).tocsc() ``` -Similarly, `Solverz.nFDAE` and `Solverz.nDAE` are respectively the numerical equation abstraction of FDAE and DAE, with the `F` and `J` attributes being the numerical interfaces. The `Solverz.nFDAE` instance has a `nstep` attribure to denote the number of historical time steps that is required. Currently, `nstep` can only be one. +Similarly, `Solverz.nFDAE` and `Solverz.nDAE` are respectively the numerical equation abstraction of FDAE and DAE, with the `F` and `J` attributes being the numerical interfaces. The `Solverz.nFDAE` instance has a `nstep` attribure to denote the number of historical time steps that is required. Currently, `nstep` can only be one. + +Sometimes, one wants to use the second derivative information. Since Release/0.1, Solverz is able to derive the Hessian-vector product, with formula + + +```{math} +\frac{\partial }{\partial y}\left(J(y)z\right)=H(y)\otimes z +``` + +where the Hessian tensor + +```{math} +H(y)=\left(\frac{\partial \nabla g_i(y)^\mathrm{T}}{y_j}\right)_{ij} +``` + +and $J(y)$ is the original Jacobian. + +One can call the `HVP()` method of numerical equations after generating numerical modules with `make_hvp=True`. A successful attempt of using `HVP` to improve the robustness of AE's solution can be found in [Solverz' cookbook](https://cook.solverz.org/ae/pf/pf.html). + #### Sparse matrix In most cases, the Jacobian is a sparse matrix in which most of the elements are zero. The sparse matrices can be decomposed very efficiently using a sparse solver. diff --git a/docs/src/reference/index.rst b/docs/src/reference/index.rst index 15f0ba6..26020b7 100644 --- a/docs/src/reference/index.rst +++ b/docs/src/reference/index.rst @@ -35,6 +35,8 @@ AE solver .. autofunction:: Solverz.solvers.nlaesolver.lm.lm +.. autofunction:: Solverz.solvers.nlaesolver.sicnm.sicnm + FDAE solver =========== diff --git a/pyproject.toml b/pyproject.toml index c844c18..e852b30 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,13 +10,13 @@ name = "Solverz" dynamic = ["version"] dependencies = [ "sympy>=1.11.1", - "numba == 0.58.1", - "numpy>=1.26.3", - "scipy>=1.12.0", + "numba >= 0.58.1", + "numpy>=2.0.0", + "scipy>=1.14.0", "dill >= 0.3.7", "pandas>=1.4.2", "openpyxl>=3.0.10", - "matplotlib>=3.5.2", + "matplotlib==3.9.0", "tqdm>=4.64.1", "pytest>=7.2.2", "pytest-datadir>=1.5.0", diff --git a/requirements.txt b/requirements.txt index 007092a..d23f44f 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,11 +1,11 @@ -numpy>=1.26.3 +numpy>=2.0.0 sympy>=1.11.1 pandas>=1.4.2 openpyxl>=3.0.10 -matplotlib>=3.5.2 +matplotlib == 3.9.0 tqdm>=4.64.1 scipy>=1.12.0 pytest>=7.2.2 networkx >= 3.1 -numba == 0.58.1 +numba >= 0.58.1 dill >= 0.3.7 diff --git a/tests/test_0.py b/tests/test_0.py index 98a6130..9c74cdb 100644 --- a/tests/test_0.py +++ b/tests/test_0.py @@ -1,4 +1,4 @@ -from Solverz import Eqn, AE, nr_method, as_Vars, iVar, made_numerical +from Solverz import Eqn, AE, nr_method, as_Vars, iVar, made_numerical, Opt x = iVar(name='x', value=2) @@ -7,7 +7,7 @@ eqn=e) y = as_Vars(x) nf = made_numerical(f, y) -sol = nr_method(nf, y) +sol = nr_method(nf, y, Opt(ite_tol=1e-8)) def test_nr_method():