diff --git a/.github/workflows/cookbook-practice.yml b/.github/workflows/cookbook-practice.yml new file mode 100644 index 0000000..6ebaf0c --- /dev/null +++ b/.github/workflows/cookbook-practice.yml @@ -0,0 +1,28 @@ +name: Run tests in cookbook + +on: [push] + +jobs: + tests_in_cookbook: + + runs-on: windows-latest + + steps: + - name: Checkout Test Repository + uses: actions/checkout@v3 + with: + repository: 'rzyu45/Solverz-Cookbook' + ref: 'main' + - uses: actions/setup-python@v4 + with: + python-version: '3.10' + cache: 'pip' # caching pip dependencies + + - name: Install Main Repository as Dependency using Git URL + run: | + pip install --upgrade pip setuptools wheel + pip install git+https://github.com/${{github.repository}}.git@${{ github.sha }} + + - name: Run Tests + run: | + pytest diff --git a/.github/workflows/doc-deploy.yml b/.github/workflows/doc-deploy.yml index a6f111a..ee0774a 100644 --- a/.github/workflows/doc-deploy.yml +++ b/.github/workflows/doc-deploy.yml @@ -19,12 +19,13 @@ jobs: with: submodules: true lfs: false - - name: Bulid Docs + - name: Deploy Python uses: actions/setup-python@v4 with: python-version: '3.11' cache: 'pip' # caching pip dependencies - - run: | + - name: Build Docs + run: | pip install -r requirements.txt cd docs pip install -r requirements.txt @@ -38,9 +39,9 @@ jobs: action: "upload" ###### Repository/Build Configurations - These values can be configured to match your app requirements. ###### # For more information regarding Static Web App workflow configurations, please visit: https://aka.ms/swaworkflowconfig - app_location: "/" # App source code path - api_location: "" # Api source code path - optional - output_location: "docs/build/html" # Built app content directory - optional + app_location: "docs/build/html" # App source code path + output_location: "" # Built app content directory - optional + skip_app_build: true ###### End of Repository/Build Configurations ###### close_pull_request_job: diff --git a/.github/workflows/publish-to-pypi.yml b/.github/workflows/publish-to-pypi.yml index f2d980f..efec1d7 100644 --- a/.github/workflows/publish-to-pypi.yml +++ b/.github/workflows/publish-to-pypi.yml @@ -1,4 +1,4 @@ -name: Publish Python 🐍 distribution 📦 to PyPI +name: Publish to Pypi on: push diff --git a/.github/workflows/python-package.yml b/.github/workflows/run-tests.yml similarity index 52% rename from .github/workflows/python-package.yml rename to .github/workflows/run-tests.yml index 90d97a4..c0b580e 100644 --- a/.github/workflows/python-package.yml +++ b/.github/workflows/run-tests.yml @@ -1,17 +1,11 @@ -# This workflow will install Python dependencies (if not cached), run tests and lint with a variety of Python versions -# For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python - -name: Python package +name: Run tests on: [push] jobs: - build: + built_in_tests: - runs-on: ${{ matrix.os }} - strategy: - matrix: - os: [ windows-latest ] + runs-on: windows-latest steps: - uses: actions/checkout@v3 diff --git a/Solverz/__init__.py b/Solverz/__init__.py index be9935d..f284f7f 100644 --- a/Solverz/__init__.py +++ b/Solverz/__init__.py @@ -6,9 +6,7 @@ from Solverz.num_api.custom_function import minmod_flag, minmod from Solverz.variable.variables import Vars, TimeVars, as_Vars from Solverz.solvers import * -from Solverz.code_printer.make_pyfunc import made_numerical -from Solverz.code_printer.py_printer import render_modules -from Solverz.code_printer.make_module import module_printer +from Solverz.code_printer import made_numerical, module_printer from Solverz.utilities.io import save, load, save_result from Solverz.utilities.profile import count_time from Solverz.variable.ssymbol import Var, AliasVar diff --git a/Solverz/code_printer/__init__.py b/Solverz/code_printer/__init__.py index e69de29..9e2ce60 100644 --- a/Solverz/code_printer/__init__.py +++ b/Solverz/code_printer/__init__.py @@ -0,0 +1,2 @@ +from Solverz.code_printer.python.inline.inline_printer import made_numerical +from Solverz.code_printer.make_module import module_printer diff --git a/Solverz/code_printer/make_module.py b/Solverz/code_printer/make_module.py index bc2a215..9bb34d4 100644 --- a/Solverz/code_printer/make_module.py +++ b/Solverz/code_printer/make_module.py @@ -1,6 +1,6 @@ from typing import List -from Solverz.code_printer.py_printer import render_modules as py_module_renderer +from Solverz.code_printer.python import py_module_renderer from Solverz.equation.equations import AE, FDAE, DAE from Solverz.variable.variables import Vars diff --git a/Solverz/code_printer/make_pyfunc.py b/Solverz/code_printer/make_pyfunc.py deleted file mode 100644 index 4c664e7..0000000 --- a/Solverz/code_printer/make_pyfunc.py +++ /dev/null @@ -1,5 +0,0 @@ -import numpy as np - -from Solverz.code_printer.py_printer import made_numerical - - diff --git a/Solverz/code_printer/test/__init__.py b/Solverz/code_printer/matlab/__init__.py similarity index 100% rename from Solverz/code_printer/test/__init__.py rename to Solverz/code_printer/matlab/__init__.py diff --git a/Solverz/code_printer/matlab_printer.py b/Solverz/code_printer/matlab/matlab_printer.py similarity index 100% rename from Solverz/code_printer/matlab_printer.py rename to Solverz/code_printer/matlab/matlab_printer.py diff --git a/Solverz/code_printer/py_printer.py b/Solverz/code_printer/py_printer.py deleted file mode 100644 index f71908e..0000000 --- a/Solverz/code_printer/py_printer.py +++ /dev/null @@ -1,893 +0,0 @@ -from __future__ import annotations - -import inspect -import warnings -from typing import Callable, Union, List, Dict -import builtins -from datetime import datetime -import re -import os - -import numpy as np -from sympy.codegen.ast import Assignment, AddAugmentedAssignment -from sympy import pycode, symbols, Function, Symbol, Number as SymNumber, Expr, Number as SymNumber -from sympy.codegen.ast import real, FunctionPrototype, FunctionDefinition, Return, FunctionCall -from sympy.utilities.lambdify import _import, _module_present, _get_namespace -from scipy.sparse import sparray -from numbers import Number - -from Solverz.equation.equations import Equations as SymEquations, FDAE as SymFDAE, DAE as SymDAE, AE as SymAE -from Solverz.variable.variables import Vars, TimeVars -from Solverz.equation.param import TimeSeriesParam -from Solverz.sym_algebra.symbols import iVar, SolDict, Para, idx, IdxSymBasic -from Solverz.sym_algebra.functions import zeros, CSC_array, Arange -from Solverz.utilities.address import Address -from Solverz.utilities.io import save -from Solverz.variable.variables import combine_Vars -from Solverz.num_api.custom_function import numerical_interface -from Solverz.num_api.num_eqn import nAE, nFDAE, nDAE - - -# %% - - -def parse_p(ae: SymEquations): - p = dict() - for param_name, param in ae.PARAM.items(): - if isinstance(param, TimeSeriesParam): - p.update({param_name: param}) - else: - p.update({param_name: param.v}) - return p - - -def parse_trigger_fun(ae: SymEquations): - func = dict() - for para_name, param in ae.PARAM.items(): - if param.trigger_fun is not None: - func.update({para_name + '_trigger_func': param.trigger_fun}) - - return func - - -def made_numerical(eqn: SymEquations, *xys, sparse=False, output_code=False): - """ - factory method of numerical equations - """ - print(f"Printing numerical codes of {eqn.name}") - eqn.assign_eqn_var_address(*xys) - code_F = print_F(eqn) - code_J = print_J(eqn, sparse) - custom_func = dict() - custom_func.update(numerical_interface) - custom_func.update(parse_trigger_fun(eqn)) - F = Solverzlambdify(code_F, 'F_', modules=[custom_func, 'numpy']) - J = Solverzlambdify(code_J, 'J_', modules=[custom_func, 'numpy']) - p = parse_p(eqn) - print('Complete!') - if isinstance(eqn, SymAE) and not isinstance(eqn, SymFDAE): - num_eqn = nAE(F, J, p) - elif isinstance(eqn, SymFDAE): - num_eqn = nFDAE(F, J, p, eqn.nstep) - elif isinstance(eqn, SymDAE): - num_eqn = nDAE(eqn.M, F, J, p) - else: - raise ValueError(f'Unknown equation type {type(eqn)}') - if output_code: - return num_eqn, {'F': code_F, 'J': code_J} - else: - return num_eqn - - -def render_modules(eqn: SymEquations, *xys, name, directory=None, numba=False): - """ - factory method of numerical equations - """ - print(f"Printing python codes of {eqn.name}...") - eqn.assign_eqn_var_address(*xys) - p = parse_p(eqn) - code_F = print_F_numba(eqn) - code_inner_F = print_inner_F(eqn) - code_sub_inner_F = print_sub_inner_F(eqn) - code_J = print_J_numba(eqn) - codes = print_inner_J(eqn, *xys) - code_inner_J = codes['code_inner_J'] - code_sub_inner_J = codes['code_sub_inner_J'] - custom_func = dict() - custom_func.update(numerical_interface) - - def print_trigger_func_code(): - code_tfuc = dict() - trigger_func = parse_trigger_fun(eqn) - # rename trigger_func - for func_name, func in trigger_func.items(): - name0 = func.__name__ - func_code = inspect.getsource(func).replace(name0, func_name) - func_code = func_code.replace('np.', '') - code_tfuc[func_name] = func_code - return code_tfuc - - code_tfunc_dict = print_trigger_func_code() - - print('Complete!') - - eqn_parameter = {} - if isinstance(eqn, SymAE) and not isinstance(eqn, SymFDAE): - eqn_type = 'AE' - elif isinstance(eqn, SymFDAE): - eqn_type = 'FDAE' - eqn_parameter.update({'nstep': eqn.nstep}) - elif isinstance(eqn, SymDAE): - eqn_type = 'DAE' - eqn_parameter.update({'M': eqn.M}) - else: - raise ValueError(f'Unknown equation type {type(eqn)}') - - if len(xys) == 1: - y = xys[0] - else: - y = xys[0] - for arg in xys[1:]: - y = combine_Vars(y, arg) - - 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} - eqn_parameter.update({'row': codes['row'], 'col': codes['col'], 'data': codes['data']}) - print(f"Rendering python modules!") - render_as_modules(name, - code_dict, - eqn_type, - p, - eqn_parameter, - y, - [custom_func, 'numpy'], - numba, - directory) - print('Complete!') - - -def is_valid_python_module_name(module_name): - pattern = re.compile(r'^[a-zA-Z_][a-zA-Z0-9_]*$') - return bool(pattern.match(module_name)) - - -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 - os.makedirs(location, exist_ok=True) - - # Create an empty __init__.py file - init_path = os.path.join(location, "__init__.py") - with open(init_path, "w") as file: - file.write(initiate_code) - - # Create the file with the dependency code - module_path = os.path.join(location, "dependency.py") - with open(module_path, "w") as file: - file.write(dependency_code) - - # Create the file with the module code - module_path = os.path.join(location, "num_func.py") - with open(module_path, "w") as file: - file.write(module_code) - - save(auxiliary, f'{location}/param_and_setting.pkl') - - -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' - code += 'import time\n' - match eqn_type: - case 'AE': - code += 'from Solverz.num_api.num_eqn import nAE\n' - code += 'mdl = nAE(F_, J_, p)\n' - code_compile = 'mdl.F(y.array, p)\nmdl.J(y.array, p)\n' - case 'FDAE': - try: - nstep = eqn_param['nstep'] - except KeyError as e: - raise ValueError("Cannot parse nstep attribute for FDAE object printing!") - code += 'from Solverz.num_api.num_eqn import nFDAE\n' - code += 'mdl = nFDAE(F_, J_, p, setting["nstep"])\n' - args_str = ', '.join(['y.array' for i in range(nstep)]) - code_compile = f'mdl.F(0, y.array, p, {args_str})\nmdl.J(0, y.array, p, {args_str})\n' - 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.array, p)\nmdl.J(0, y.array, p)\n' - case _: - raise ValueError(f'Unknown equation type {eqn_type}') - code += f'print("Compiling model {module_name}...")\n' - code += f'start = time.perf_counter()\n' - code += code_compile - code += f'end = time.perf_counter()\n' - code += "print(f'Compiling time elapsed: {end-start}s')\n" - return code - - -def print_module_code(code_dict: Dict[str, str], numba=False): - code = 'from .dependency import *\n' - code += """_data_ = setting["data"]\n""" - code += """_F_ = zeros_like(y, dtype=float64)""" - code += '\n\r\n' - for tfunc in code_dict['tfunc_dict'].values(): - if numba: - 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: - 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: - if numba: - code += '@njit(cache=True)\n' - code += sub_func - code += '\n\r\n' - return code - - -def print_dependency_code(modules): - code = "import os\n" - code += "current_module_dir = os.path.dirname(os.path.abspath(__file__))\n" - code += 'from Solverz import load\n' - code += 'auxiliary = load(f"{current_module_dir}\\\\param_and_setting.pkl")\n' - code += 'from numpy import *\n' - code += 'from numpy import abs\n' - code += 'from Solverz.num_api.custom_function import *\n' # import Solverz built-in func - code += 'from scipy.sparse import *\n' - code += 'from numba import njit\n' - code += 'setting = auxiliary["eqn_param"]\n' - code += 'row = setting["row"]\n' - code += 'col = setting["col"]\n' - code += 'data_ = setting["data"]\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): - if is_valid_python_module_name(name): - module_name = name - else: - raise ValueError("Invalid python module name!") - - init_code = print_init_code(eqn_type, module_name, eqn_param) - module_code = print_module_code(code_dict, numba=numba) - dependency_code = print_dependency_code(modules) - - create_python_module(module_name, - init_code, - module_code, - dependency_code, - {'p': p, 'eqn_param': eqn_param, 'vars': variables}, - directory) - - -def parse_name_space(modules=None): - # If the user hasn't specified any modules, use what is available. - if modules is None: - try: - _import("scipy") - except ImportError: - try: - _import("numpy") - except ImportError: - # Use either numpy (if available) or python.math where possible. - # XXX: This leads to different behaviour on different systems and - # might be the reason for irreproducible errors. - modules = ["math", "mpmath", "sympy"] - else: - modules = ["numpy"] - else: - modules = ["numpy", "scipy"] - - # Get the needed namespaces. - namespaces = [] - - # Check for dict before iterating - if isinstance(modules, (dict, str)) or not hasattr(modules, '__iter__'): - namespaces.append(modules) - else: - # consistency check - if _module_present('numexpr', modules) and len(modules) > 1: - raise TypeError("numexpr must be the only item in 'modules'") - namespaces += list(modules) - # fill namespace with first having highest priority - namespace = {} - for m in namespaces[::-1]: - buf = _get_namespace(m) - namespace.update(buf) - - # Provide lambda expression with builtins, and compatible implementation of range - namespace.update({'builtins': builtins, 'range': range}) - return namespace - - -def Solverzlambdify(funcstr, funcname, modules=None): - """Convert a Solverz numerical f/g/F/J evaluation expression into a function that allows for fast - numeric evaluation. - """ - namespace = parse_name_space(modules) - - funclocals = {} - current_time = datetime.now() - filename = '' % current_time - c = compile(funcstr, filename, 'exec') - exec(c, namespace, funclocals) - - func = funclocals[funcname] - # Apply the docstring - src_str = funcstr - # TODO: should collect and show the module imports from the code printers instead of the namespace - func.__doc__ = ( - "Created with Solverz at \n\n" - "{sig}\n\n" - "Source code:\n\n" - "{src}\n\n" - ).format(sig=current_time, src=src_str) - return func - - -def _print_var_parser(var_name, suffix: str, var_address): - y = iVar('y_' + suffix, internal_use=True) - if suffix == '': - return Assignment(iVar(var_name + suffix), y[var_address]) - else: - return Assignment(iVar(var_name + '_tag_' + suffix), y[var_address]) - pass - - -def print_var(ae: SymEquations, numba_printer=False): - var_declaration = [] - if hasattr(ae, 'nstep'): - nstep = ae.nstep - else: - nstep = 0 - suffix = [''] + [f'{i}' for i in range(nstep)] - var_list = [] - for i in suffix: - for var_name in ae.var_address.v.keys(): - var_address = ae.var_address[var_name] - var_assign = _print_var_parser(var_name, i, var_address) - var_declaration.append(var_assign) - var_list.append(var_assign.lhs) - if not numba_printer: - return var_declaration - else: - return var_declaration, var_list - - -def print_param(ae: SymEquations, numba_printer=False): - param_declaration = [] - p = SolDict('p_') - param_list = [] - for symbol_name, symbol in ae.SYMBOLS.items(): - if isinstance(symbol, (Para, idx)): - if isinstance(ae.PARAM[symbol_name], TimeSeriesParam): - param_assign = Assignment(Para(symbol_name), FunctionCall(f'p_["{symbol_name}"].get_v_t', 't')) - else: - param_assign = Assignment(Para(symbol_name), p[symbol_name]) - param_declaration.append(param_assign) - param_list.append(param_assign.lhs) - if not numba_printer: - return param_declaration - else: - return param_declaration, param_list - - -def _print_F_assignment(eqn_address, eqn, i): - F_ = iVar('_F_', internal_use=True) - return Assignment(F_[eqn_address], FunctionCall(f'inner_F{int(i)}', list(eqn.SYMBOLS.values()))) - - -def print_eqn_assignment(ae: SymEquations, numba_printer=False): - _F_ = iVar('_F_', internal_use=True) - eqn_declaration = [] - if not numba_printer: - eqn_declaration.append(Assignment(_F_, zeros(ae.eqn_size, ))) - count = 0 - for eqn_name in ae.EQNs.keys(): - eqn_address = ae.a[eqn_name] - eqn = ae.EQNs[eqn_name] - if numba_printer: - 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}") - _F_ = iVar('_F_', internal_use=True) - eqn_declaration.append(Assignment(_F_[eqn_address], - FunctionCall(f'inner_F{int(count)}', list(eqn.SYMBOLS.values())))) - count = count + 1 - else: - eqn_declaration.append(Assignment(_F_[eqn_address], eqn.RHS)) - return eqn_declaration - - -def _parse_jac_eqn_address(eqn_address: slice, derivative_dim, sparse): - if eqn_address.stop - eqn_address.start == 1: - if derivative_dim == 1 or derivative_dim == 0: - eqn_address = eqn_address.start if not sparse else SolList(eqn_address.start) - else: - if sparse: - return iVar(f'arange({eqn_address.start}, {eqn_address.stop})')[idx('value_coo.row')] - else: - if derivative_dim == 1 or derivative_dim == 0: - eqn_address = Arange(eqn_address.start, eqn_address.stop) - else: - if sparse: - return iVar(f'arange({eqn_address.start}, {eqn_address.stop})')[idx('value_coo.row')] - return eqn_address - - -def _parse_jac_var_address(var_address_slice: slice, - derivative_dim, - var_idx, - sparse, - eqn_length: int = 1): - if var_idx is not None: - try: # try to simplify the variable address in cases such as [1:10][0,1,2,3] - temp = np.arange(var_address_slice.start, var_address_slice.stop)[var_idx] - if isinstance(temp, (Number, SymNumber)): # number a - if not sparse: - var_address = temp - else: - if derivative_dim == 0: # For the derivative of omega[0] in eqn omega-omega[0] where omega is a vector - var_address = eqn_length * SolList(temp) - elif derivative_dim == 1: - var_address = SolList(temp) - else: - var_address = iVar('array([' + str(temp) + '])')[idx('value_coo.col')] - else: # np.ndarray - if np.all(np.diff(temp) == 1): # list such as [1,2,3,4] can be viewed as slice [1:5] - if derivative_dim == 2: - if sparse: - return iVar(f'arange({temp[0]}, {temp[-1] + 1})')[idx('value_coo.col')] - else: - var_address = slice(temp[0], temp[-1] + 1) - elif derivative_dim == 1 or derivative_dim == 0: - var_address = Arange(temp[0], temp[-1] + 1) - else: # arbitrary list such as [1,3,4,5,6,9] - if not sparse or derivative_dim < 2: - var_address = SolList(*temp) - else: - var_address = iVar('array([' + ','.join([str(ele) for ele in temp]) + '])')[idx('value_coo.col')] - except (TypeError, IndexError): - if isinstance(var_idx, str): - var_idx = idx(var_idx) - if not sparse or derivative_dim < 2: - var_address = iVar(f"arange({var_address_slice.start}, {var_address_slice.stop})")[var_idx] - else: - var_address = iVar(f"arange({var_address_slice.start}, {var_address_slice.stop})")[ - var_idx[idx('value_coo.col')]] - else: - if derivative_dim == 2: - if sparse: - return iVar(f'arange({var_address_slice.start}, {var_address_slice.stop})')[idx('value_coo.col')] - else: - var_address = var_address_slice - elif derivative_dim == 1 or derivative_dim == 0: - var_address = Arange(var_address_slice.start, var_address_slice.stop) - return var_address - - -def _parse_jac_data(data_length, derivative_dim: int, rhs: Union[Expr, Number, SymNumber], rhs_v_type='array'): - if derivative_dim == 2: - return iVar('value_coo.data') - elif derivative_dim == 1: - return rhs - elif derivative_dim == 0: - if rhs_v_type == 'Number': # if rhs is a number, then return length*[rhs] - return data_length * SolList(rhs) - elif rhs_v_type == 'array': # if rhs produces np.ndarray then return length*rhs.tolist() - return data_length * tolist(rhs) - - -def print_J_block(eqn_address_slice, var_address_slice, derivative_dim, var_idx, rhs, sparse, - rhs_v_dtpe='array') -> List: - if sparse: - eqn_address = _parse_jac_eqn_address(eqn_address_slice, - derivative_dim, - True) - var_address = _parse_jac_var_address(var_address_slice, - derivative_dim, - var_idx, - True, - eqn_address_slice.stop - eqn_address_slice.start) - # assign elements to sparse matrix can not be easily broadcast, so we have to parse the data - data = _parse_jac_data(eqn_address_slice.stop - eqn_address_slice.start, - derivative_dim, - rhs, - rhs_v_dtpe) - if derivative_dim < 2: - return [extend(iVar('row', internal_use=True), eqn_address), - extend(iVar('col', internal_use=True), var_address), - extend(iVar('data', internal_use=True), data)] - else: - return [Assignment(iVar('value_coo'), coo_array(rhs)), - extend(iVar('row', internal_use=True), eqn_address), - extend(iVar('col', internal_use=True), var_address), - extend(iVar('data', internal_use=True), data)] - else: - eqn_address = _parse_jac_eqn_address(eqn_address_slice, - derivative_dim, - False) - var_address = _parse_jac_var_address(var_address_slice, - derivative_dim, - var_idx, - False) - return [AddAugmentedAssignment(iVar('J_', internal_use=True)[eqn_address, var_address], rhs)] - - -def print_J_dense(ae: SymEquations): - eqn_declaration = [] - for eqn_name, eqn in ae.EQNs.items(): - eqn_address_slice = ae.a[eqn_name] - for var_name, eqndiff in eqn.derivatives.items(): - derivative_dim = eqndiff.dim - if derivative_dim < 0: - raise ValueError("Derivative dimension not assigned") - var_address_slice = ae.var_address[eqndiff.diff_var_name] - var_idx = eqndiff.var_idx - rhs = eqndiff.RHS - eqn_declaration.extend(print_J_block(eqn_address_slice, - var_address_slice, - derivative_dim, - var_idx, - rhs, - False)) - return eqn_declaration - - -def print_J_sparse(ae: SymEquations): - eqn_declaration = [] - for eqn_name, eqn in ae.EQNs.items(): - eqn_address_slice = ae.a[eqn_name] - for var_name, eqndiff in eqn.derivatives.items(): - derivative_dim = eqndiff.dim - if derivative_dim < 0: - raise ValueError("Derivative dimension not assigned") - var_address_slice = ae.var_address[eqndiff.diff_var_name] - var_idx = eqndiff.var_idx - rhs = eqndiff.RHS - rhs_v_type = eqndiff.v_type - eqn_declaration.extend(print_J_block(eqn_address_slice, - var_address_slice, - derivative_dim, - var_idx, - rhs, - True, - rhs_v_type)) - return eqn_declaration - - -def _print_trigger_func(para_name, trigger_var: List[str]): - return Assignment(Para(para_name), - FunctionCall(para_name + '_trigger_func', tuple([symbols(var) for var in trigger_var]))) - - -def print_trigger(ae: SymEquations): - trigger_declaration = [] - for name, param in ae.PARAM.items(): - if param.triggerable: - trigger_declaration.append(_print_trigger_func(name, - param.trigger_var)) - return trigger_declaration - - -def print_func_prototype(ae: SymEquations, func_name: str): - t, y_, p_ = symbols('t y_ p_', real=True) - if isinstance(ae, (SymDAE, SymFDAE)): - args = [t, y_, p_] - else: - args = [y_, p_] - xtra_args = [] - if hasattr(ae, 'nstep'): - if ae.nstep > 0: - xtra_args.extend([symbols('y_' + f'{i}', real=True) for i in range(ae.nstep)]) - - fp = FunctionPrototype(real, func_name, args + xtra_args) - return fp - - -def print_F(ae: SymEquations): - fp = print_func_prototype(ae, 'F_') - body = [] - body.extend(print_var(ae)) - body.extend(print_param(ae)) - body.extend(print_trigger(ae)) - body.extend(print_eqn_assignment(ae)) - temp = iVar('_F_', internal_use=True) - body.extend([Return(temp)]) - fd = FunctionDefinition.from_FunctionPrototype(fp, body) - return pycode(fd, fully_qualified_modules=False) - - -def print_J(ae: SymEquations, sparse=False): - fp = print_func_prototype(ae, 'J_') - # initialize temp - temp = iVar('J_', internal_use=True) - body = list() - body.extend(print_var(ae)) - body.extend(print_param(ae)) - body.extend(print_trigger(ae)) - if not sparse: - body.append(Assignment(temp, zeros(ae.eqn_size, ae.vsize))) - body.extend(print_J_dense(ae)) - 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_sparse(ae)) - body.append(Return(coo_2_csc(ae))) - Jd = FunctionDefinition.from_FunctionPrototype(fp, body) - return pycode(Jd, fully_qualified_modules=False) - - -def print_J_numba(ae: SymEquations): - fp = print_func_prototype(ae, 'J_') - body = [] - var_assignments, var_list = print_var(ae, numba_printer=True) - body.extend(var_assignments) - param_assignments, param_list = print_param(ae, numba_printer=True) - body.extend(param_assignments) - body.extend(print_trigger(ae)) - body.extend([Assignment(iVar('data', internal_use=True), - FunctionCall('inner_J', [symbols('_data_', real=True)] + var_list + param_list))]) - body.extend([Return(coo_2_csc(ae))]) - fd = FunctionDefinition.from_FunctionPrototype(fp, body) - return pycode(fd, fully_qualified_modules=False) - - -def print_inner_J(ae: SymEquations, *xys): - var_assignments, var_list = print_var(ae, numba_printer=True) - param_assignments, param_list = print_param(ae, numba_printer=True) - args = [] - for var in var_list + param_list: - args.append(symbols(var.name, real=True)) - fp = FunctionPrototype(real, 'inner_J', [symbols('_data_', real=True)] + args) - body = [] - row, col, jac_address = parse_jac_address(ae, *xys) - data = np.zeros_like(row, dtype=np.float64) - # temp = iVar('data_', internal_use=True) - # body.extend([Assignment(temp, zeros(jac_address.total_size, ))]) - code_sub_inner_J_blocks = [] - count = 0 - for jac_ in jac_address.object_list: - eqn_name, var_name = jac_.split("@@@") - derivative = ae.EQNs[eqn_name].derivatives[var_name] - rhs = derivative.RHS - if isinstance(rhs, (Number, SymNumber)): - data[jac_address[jac_]] = rhs - else: - body.append(Assignment(iVar('_data_', internal_use=True)[jac_address[jac_]], - FunctionCall(f'inner_J{int(count)}', list(derivative.SYMBOLS.values())))) - args1 = [] - for var in derivative.SYMBOLS.keys(): - args1.append(symbols(var, real=True)) - - fp1 = FunctionPrototype(real, f'inner_J{int(count)}', args1) - body1 = [Return(rhs)] - fd1 = FunctionDefinition.from_FunctionPrototype(fp1, body1) - code_sub_inner_J_blocks.append(pycode(fd1, fully_qualified_modules=False)) - count += 1 - temp = iVar('_data_', internal_use=True) - body.extend([Return(temp)]) - fd = FunctionDefinition.from_FunctionPrototype(fp, body) - return {'code_inner_J': pycode(fd, fully_qualified_modules=False), - 'code_sub_inner_J': code_sub_inner_J_blocks, - 'row': row, - 'col': col, - 'data': data} - - -def _print_inner_J_block_assignment(rhs, address: slice): - data = iVar('data_', internal_use=True) - return Assignment(data[address], rhs) - - -def print_F_numba(ae: SymEquations): - fp = print_func_prototype(ae, 'F_') - body = [] - var_assignments, var_list = print_var(ae, numba_printer=True) - body.extend(var_assignments) - param_assignments, param_list = print_param(ae, numba_printer=True) - body.extend(param_assignments) - body.extend(print_trigger(ae)) - body.extend( - [Return(FunctionCall('inner_F', [symbols('_F_', real=True)] + var_list + param_list))]) - fd = FunctionDefinition.from_FunctionPrototype(fp, body) - return pycode(fd, fully_qualified_modules=False) - - -def print_inner_F(ae: SymEquations): - var_assignments, var_list = print_var(ae, numba_printer=True) - param_assignments, param_list = print_param(ae, numba_printer=True) - args = [] - for var in var_list + param_list: - args.append(symbols(var.name, real=True)) - fp = FunctionPrototype(real, 'inner_F', [symbols('_F_', real=True)] + args) - body = [] - body.extend(print_eqn_assignment(ae, True)) - temp = iVar('_F_', internal_use=True) - body.extend([Return(temp)]) - fd = FunctionDefinition.from_FunctionPrototype(fp, body) - return pycode(fd, fully_qualified_modules=False) - - -def print_sub_inner_F(ae: SymEquations): - code_blocks = [] - count = 0 - for eqn_name in ae.EQNs.keys(): - eqn = ae.EQNs[eqn_name] - args = [] - for var in eqn.SYMBOLS.keys(): - args.append(symbols(var, real=True)) - fp = FunctionPrototype(real, f'inner_F{count}', args) - body = [Return(eqn.RHS)] - fd = FunctionDefinition.from_FunctionPrototype(fp, body) - count = count + 1 - code_blocks.append(pycode(fd, fully_qualified_modules=False)) - return code_blocks - - -def parse_jac_address(eqns: SymEquations, *xys): - if isinstance(eqns, SymAE): - gy = eqns.g_y(*xys) - elif isinstance(eqns, SymDAE): - gy = eqns.f_xy(0, *xys) - if len(eqns.g_list) > 0: - gy.extend(eqns.g_xy(0, *xys)) - else: - raise NotImplementedError(f"Unknown equation type {type(eqns)}") - - row = np.array([], dtype=int) - col = np.array([], dtype=int) - jac_block_address = Address() - for gy_tuple in gy: - eqn_name = gy_tuple[0] - # eqn_diffs: Dict[str, EqnDiff] = eqns.EQNs[eqn_name].derivatives - var_name = gy_tuple[1] - eqndiff = gy_tuple[2] - diff_var = eqndiff.diff_var - diff_var_name = eqndiff.diff_var.name - value = gy_tuple[3] - eqn_address = eqns.a[eqn_name] - var_address = eqns.var_address[var_name] - if isinstance(value, (np.ndarray, sparray)): - if value.ndim == 2: # matrix - raise TypeError("Two-dimensional array not applicable for numba printer!\n Try rewrite the Equations!") - elif value.ndim == 1 and value.shape[0] != 1: # vector - num_jac_element = value.shape[0] - if num_jac_element != eqn_address.stop - eqn_address.start: - raise ValueError("Number of jac block elements not compatible with equation length!") - elif value.ndim == 1 and value.shape[0] == 1: # scalar in np.ndarray for example array([0.0]) - num_jac_element = eqn_address.stop - eqn_address.start - else: - raise ValueError("Unknown derivative value dimension type!") - elif isinstance(value, (Number, SymNumber)): - num_jac_element = eqn_address.stop - eqn_address.start - else: - raise ValueError(f"Unknown derivative data type {type(value)}!") - eqn_address_range = np.arange(eqn_address.start, eqn_address.stop) - row = np.append(row, eqn_address_range) - if isinstance(diff_var, IdxSymBasic): - index = diff_var.index - if not isinstance(index, (int, slice, np.integer)): - raise TypeError( - f"Index type {type(diff_var.index)} not applicable for numba printer!\n Try rewrite the Variable!") - else: - # reshape is to convert float/integer to 1-dim numpy.ndarray - var_address_range = np.array(np.arange(var_address.start, var_address.stop)[index]).reshape((-1,)) - if len(var_address_range) < len(eqn_address_range): - warnings.warn(f'Address of variable {diff_var_name} (length={len(var_address_range)}) shorter than equation address of {eqn_name} (length={len(eqn_address_range)}). Please check the variable address and equation address of this part.') - var_address_range = np.array(len(eqn_address_range)*(var_address_range.tolist())) - elif isinstance(diff_var, iVar): - var_address_range = np.arange(var_address.start, var_address.stop) - if len(var_address_range) != len(eqn_address_range): - raise ValueError('Equation address range is different from variable address range') - col = np.append(col, var_address_range) - jac_block_address.add(f'{eqn_name}@@@{diff_var_name}', num_jac_element) - - assert len(row) == len(col) == jac_block_address.total_size - - return row, col, jac_block_address - - -class coo_2_csc(Symbol): - - def __new__(cls, ae: SymEquations): - obj = Symbol.__new__(cls, f'coo_2_csc: {ae.name}') - obj.eqn_size = ae.eqn_size - obj.vsize = ae.vsize - return obj - - def _numpycode(self, printer, **kwargs): - return f'coo_array((data, (row,col)), ({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.") - - def _numpycode(self, printer, **kwargs): - return f'coo_array({printer._print(self.args[0])})' - - def _pythoncode(self, printer, **kwargs): - return self._numpycode(printer, **kwargs) - - -class extend(Function): - - def _numpycode(self, printer, **kwargs): - return f'{printer._print(self.args[0])}.extend({printer._print(self.args[1])})' - - def _pythoncode(self, printer, **kwargs): - return self._numpycode(printer, **kwargs) - - -class SolList(Function): - - # @classmethod - # def eval(cls, *args): - # if any([not isinstance(arg, Number) for arg in args]): - # raise ValueError(f"Solverz' list object accepts only number inputs.") - - def _numpycode(self, printer, **kwargs): - return r'[' + ', '.join([printer._print(arg) for arg in self.args]) + r']' - - def _pythoncode(self, printer, **kwargs): - return self._numpycode(printer, **kwargs) - - -class Array(Function): - - def _numpycode(self, printer, **kwargs): - return r'array([' + ', '.join([printer._print(arg) for arg in self.args]) + r'])' - - def _pythoncode(self, printer, **kwargs): - return self._numpycode(printer, **kwargs) - - -class tolist(Function): - - @classmethod - def eval(cls, *args): - if len(args) != 1: - raise ValueError(f"Solverz' tolist function accepts only one input.") - - def _numpycode(self, printer, **kwargs): - return r'((' + printer._print(self.args[0]) + r').tolist())' - - def _pythoncode(self, printer, **kwargs): - return self._numpycode(printer, **kwargs) diff --git a/Solverz/code_printer/python/__init__.py b/Solverz/code_printer/python/__init__.py new file mode 100644 index 0000000..094f9d3 --- /dev/null +++ b/Solverz/code_printer/python/__init__.py @@ -0,0 +1 @@ +from Solverz.code_printer.python.module.module_generator import render_modules as py_module_renderer diff --git a/Solverz/code_printer/python/inline/__init__.py b/Solverz/code_printer/python/inline/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/Solverz/code_printer/python/inline/inline_printer.py b/Solverz/code_printer/python/inline/inline_printer.py new file mode 100644 index 0000000..a9b977a --- /dev/null +++ b/Solverz/code_printer/python/inline/inline_printer.py @@ -0,0 +1,225 @@ +from __future__ import annotations + +from sympy.utilities.lambdify import _import, _module_present, _get_namespace + +from Solverz.variable.variables import Vars +from Solverz.code_printer.python.utilities import * + + +# %% + + +def print_F(eqs_type: str, + EQNs: Dict[str, Eqn], + EqnAddr: Address, + var_addr: Address, + PARAM: Dict[str, ParamBase], + nstep: int = 0): + fp = print_F_J_prototype(eqs_type, + 'F_', + nstep) + body = [] + body.extend(print_var(var_addr, + nstep)[0]) + param_decla = print_param(PARAM)[0] + body.extend(param_decla) + body.extend(print_trigger(PARAM)) + body.extend(print_eqn_assignment(EQNs, + EqnAddr, + False)) + temp = iVar('_F_', internal_use=True) + body.extend([Return(temp)]) + fd = FunctionDefinition.from_FunctionPrototype(fp, body) + return pycode(fd, fully_qualified_modules=False) + + +def print_J(eqs_type: str, + jac: Jac, + EqnAddr: Address, + var_addr: Address, + PARAM: Dict[str, ParamBase], + nstep: int = 0, + sparse=False): + fp = print_F_J_prototype(eqs_type, + 'J_', + nstep) + # initialize temp + temp = iVar('J_', 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(jac, 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(jac, 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 print_J_blocks(jac: Jac, sparse: bool): + eqn_declaration = [] + for eqn_name, jbs_row in jac.blocks.items(): + for var, jb in jbs_row.items(): + eqn_declaration.extend(print_J_block(jb, + sparse)) + return eqn_declaration + + +def print_J_block(jb: JacBlock, sparse: bool) -> List: + if sparse: + match jb.DeriType: + case 'matrix': + # return [Assignment(iVar('value_coo'), coo_array(rhs)), + # 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!") + 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('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): + """ + factory method of numerical equations + """ + print(f"Printing numerical codes of {eqs.name}") + eqs.FormJac(y) + code_F = print_F(eqs.__class__.__name__, + eqs.EQNs, + eqs.a, + eqs.var_address, + eqs.PARAM, + eqs.nstep) + code_J = print_J(eqs.__class__.__name__, + eqs.jac, + eqs.a, + eqs.var_address, + eqs.PARAM, + eqs.nstep, + sparse) + 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']) + p = parse_p(eqs.PARAM) + print('Complete!') + if isinstance(eqs, SymAE) and not isinstance(eqs, SymFDAE): + num_eqn = nAE(F, J, p) + elif isinstance(eqs, SymFDAE): + num_eqn = nFDAE(F, J, p, eqs.nstep) + elif isinstance(eqs, SymDAE): + num_eqn = nDAE(eqs.M, F, J, p) + else: + raise ValueError(f'Unknown equation type {type(eqs)}') + if output_code: + return num_eqn, {'F': code_F, 'J': code_J} + else: + return num_eqn + + +def parse_name_space(modules=None): + # If the user hasn't specified any modules, use what is available. + if modules is None: + try: + _import("scipy") + except ImportError: + try: + _import("numpy") + except ImportError: + # Use either numpy (if available) or python.math where possible. + # XXX: This leads to different behaviour on different systems and + # might be the reason for irreproducible errors. + modules = ["math", "mpmath", "sympy"] + else: + modules = ["numpy"] + else: + modules = ["numpy", "scipy"] + + # Get the needed namespaces. + namespaces = [] + + # Check for dict before iterating + if isinstance(modules, (dict, str)) or not hasattr(modules, '__iter__'): + namespaces.append(modules) + else: + # consistency check + if _module_present('numexpr', modules) and len(modules) > 1: + raise TypeError("numexpr must be the only item in 'modules'") + namespaces += list(modules) + # fill namespace with first having highest priority + namespace = {} + for m in namespaces[::-1]: + buf = _get_namespace(m) + namespace.update(buf) + + # Provide lambda expression with builtins, and compatible implementation of range + namespace.update({'builtins': builtins, 'range': range}) + return namespace + + +def Solverzlambdify(funcstr, funcname, modules=None): + """Convert a Solverz numerical f/g/F/J evaluation expression into a function that allows for fast + numeric evaluation. + """ + namespace = parse_name_space(modules) + + funclocals = {} + current_time = datetime.now() + filename = '' % current_time + c = compile(funcstr, filename, 'exec') + exec(c, namespace, funclocals) + + func = funclocals[funcname] + # Apply the docstring + src_str = funcstr + # TODO: should collect and show the module imports from the code printers instead of the namespace + func.__doc__ = ( + "Created with Solverz at \n\n" + "{sig}\n\n" + "Source code:\n\n" + "{src}\n\n" + ).format(sig=current_time, src=src_str) + return func + + +class SolList(Function): + + # @classmethod + # def eval(cls, *args): + # if any([not isinstance(arg, Number) for arg in args]): + # raise ValueError(f"Solverz' list object accepts only number inputs.") + + def _numpycode(self, printer, **kwargs): + return r'[' + ', '.join([printer._print(arg) for arg in self.args]) + r']' + + def _pythoncode(self, printer, **kwargs): + return self._numpycode(printer, **kwargs) + + +class tolist(Function): + + @classmethod + def eval(cls, *args): + if len(args) != 1: + raise ValueError(f"Solverz' tolist function accepts only one input.") + + def _numpycode(self, printer, **kwargs): + return r'((' + printer._print(self.args[0]) + r').tolist())' + + def _pythoncode(self, printer, **kwargs): + return self._numpycode(printer, **kwargs) diff --git a/Solverz/code_printer/python/inline/tests/__init__.py b/Solverz/code_printer/python/inline/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/Solverz/code_printer/python/inline/tests/test_inline_printer.py b/Solverz/code_printer/python/inline/tests/test_inline_printer.py new file mode 100644 index 0000000..53236d7 --- /dev/null +++ b/Solverz/code_printer/python/inline/tests/test_inline_printer.py @@ -0,0 +1,267 @@ +from sympy import pycode +from numbers import Number + +import numpy as np +from numpy.testing import assert_allclose +import re +import pytest +from sympy import sin +from sympy.codegen.ast import Assignment, AddAugmentedAssignment + +from Solverz import Model, Var, AliasVar +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.variable.variables import combine_Vars, as_Vars +from Solverz.equation.jac import JacBlock, Ones, Jac +from Solverz.sym_algebra.functions import Diag +from Solverz.code_printer.python.inline.inline_printer import print_J_block, extend, SolList, print_J_blocks, print_J, \ + print_F, made_numerical + +# %% +row = iVar('row', internal_use=True) +col = iVar('col', internal_use=True) +data = iVar('data', internal_use=True) +J_ = iVar('J_', internal_use=True) + + +def test_jb_printer_scalar_var_scalar_deri(): + # non-index var + jb = JacBlock('a', + slice(0, 3), + iVar('x'), + np.array([1]), + slice(1, 2), + iVar('y'), + np.array([1])) + + symJb = print_J_block(jb, True) + assert symJb[0] == extend(row, SolList(0, 1, 2)) + 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)) + + +def test_jb_printer_vector_var_vector_deri(): + jb = JacBlock('a', + slice(1, 10), + iVar('x'), + np.ones(9), + slice(0, 9), + iVar('y'), + np.ones(9)) + symJb = print_J_block(jb, True) + assert symJb[0] == extend(row, SolList(*np.arange(1, 10).tolist())) + 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'))) + + +def test_jbs_printer(): + jac = Jac() + x = iVar('x') + y = iVar('y') + omega = iVar('omega') + jac.add_block('a', + x[0], + JacBlock('a', + slice(0, 3), + x[0], + np.array([1]), + slice(1, 2), + iVar('y'), + np.array([1]))) + jac.add_block('b', + omega[1:4], + JacBlock('b', + slice(3, 12), + omega[4:13], + np.ones(9), + slice(3, 100), + iVar('y') ** 2, + np.ones(9))) + symJbs = print_J_blocks(jac, True) + assert symJbs[0] == extend(row, SolList(0, 1, 2)) + assert symJbs[1] == extend(col, SolList(1, 1, 1)) + assert symJbs[2] == extend(data, y * Ones(3)) + assert symJbs[3] == extend(row, SolList(3, 4, 5, 6, 7, 8, 9, 10, 11)) + 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[1] == AddAugmentedAssignment(J_[3:12, 7:16], Diag(y ** 2)) + + +expected_J_den = """def J_(t, y_, p_): + h = y_[0:1] + v = y_[1:2] + g = p_["g"] + J_ = zeros((2, 2)) + J_[0:1,1:2] += ones(1) + return J_ +""".strip() + +expected_J_sp = """def J_(t, y_, p_): + h = y_[0:1] + v = y_[1:2] + g = p_["g"] + row = [] + col = [] + data = [] + row.extend([0]) + col.extend([1]) + data.extend(ones(1)) + return coo_array((data, (row, col)), (2, 2)).tocsc() +""".strip() + +expected_F = """def F_(t, y_, p_): + h = y_[0:1] + v = y_[1:2] + g = p_["g"] + _F_ = zeros((2, )) + _F_[0:1] = v + _F_[1:2] = g + return _F_ +""".strip() + + +def test_print_F_J(): + m = Model() + m.h = Var('h', 0) + m.v = Var('v', 20) + m.g = Param('g', -9.8) + m.f1 = Ode('f1', f=m.v, diff_var=m.h) + m.f2 = Ode('f2', f=m.g, diff_var=m.v) + bball, y0 = m.create_instance() + assert print_J(bball.__class__.__name__, + bball.jac, + bball.a, + bball.var_address, + bball.PARAM, + bball.nstep) == expected_J_den + assert print_J(bball.__class__.__name__, + bball.jac, + bball.a, + bball.var_address, + bball.PARAM, + bball.nstep, + True) == expected_J_sp + assert print_F(bball.__class__.__name__, + bball.EQNs, + bball.a, + bball.var_address, + bball.PARAM, + bball.nstep) == expected_F + nbball = made_numerical(bball, y0, sparse=True) + np.testing.assert_allclose(nbball.F(0, y0, nbball.p), + np.array([20, -9.8])) + np.testing.assert_allclose(nbball.J(0, y0, nbball.p).toarray(), + np.array([[0., 1.], [0., 0.]])) + nbball = made_numerical(bball, y0, sparse=False) + np.testing.assert_allclose(nbball.J(0, y0, nbball.p), + np.array([[0., 1.], [0., 0.]])) + + +expected_J_den_fdae = """def J_(t, y_, p_, y_0): + p = y_[0:3] + q = y_[3:6] + p_tag_0 = y_0[0:3] + q_tag_0 = y_0[3:6] + 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) + return J_ +""".strip() + +expected_J_sp_fdae = """def J_(t, y_, p_, y_0): + p = y_[0:3] + q = y_[3:6] + p_tag_0 = y_0[0:3] + q_tag_0 = y_0[3:6] + row = [] + col = [] + data = [] + row.extend([0, 1]) + col.extend([1, 2]) + data.extend(ones(2)) + row.extend([2, 3]) + col.extend([0, 1]) + data.extend(-ones(2)) + row.extend([4]) + col.extend([2]) + data.extend(-ones(1)) + row.extend([5]) + col.extend([2]) + data.extend(ones(1)) + return coo_array((data, (row, col)), (6, 6)).tocsc() +""".strip() + +expected_F_fdae = """def F_(t, y_, p_, y_0): + p = y_[0:3] + q = y_[3:6] + p_tag_0 = y_0[0:3] + q_tag_0 = y_0[3:6] + _F_ = zeros((6, )) + _F_[0:2] = -p_tag_0[0:2] + p[1:3] + _F_[2:4] = p_tag_0[1:3] - p[0:2] + _F_[4:5] = p_tag_0[0] - p[2] + _F_[5:6] = -p_tag_0[0] + p[2] + return _F_ +""".strip() + + +def test_print_F_J_FDAE(): + m = Model() + m.p = Var('p', value=[2, 3, 4]) + m.q = Var('q', value=[0, 1, 2]) + m.p0 = AliasVar('p', init=m.p) + m.q0 = AliasVar('q', init=m.q) + + m.ae1 = Eqn('cha1', + m.p[1:3] - m.p0[0:2]) + + m.ae2 = Eqn('cha2', + m.p0[1:3] - m.p[0:2]) + m.ae3 = Eqn('cha3', + m.p0[0] - m.p[2]) + m.ae4 = Eqn('cha4', + m.p[2] - m.p0[0]) + fdae, y0 = m.create_instance() + assert print_J(fdae.__class__.__name__, + fdae.jac, + fdae.a, + fdae.var_address, + fdae.PARAM, + fdae.nstep) == expected_J_den_fdae + assert print_J(fdae.__class__.__name__, + fdae.jac, + fdae.a, + fdae.var_address, + fdae.PARAM, + fdae.nstep, + True) == expected_J_sp_fdae + assert print_F(fdae.__class__.__name__, + fdae.EQNs, + fdae.a, + fdae.var_address, + fdae.PARAM, + fdae.nstep) == expected_F_fdae + + +def test_made_numerical(): + x = iVar('x', [1, 1]) + f1 = Eqn('f1', 2 * x[0] + x[1]) + f2 = Eqn('f2', x[0] ** 2 + sin(x[1])) + + F = AE([f1, f2]) + y = as_Vars([x]) + 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) diff --git a/Solverz/code_printer/python/module/__init__.py b/Solverz/code_printer/python/module/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/Solverz/code_printer/python/module/module_generator.py b/Solverz/code_printer/python/module/module_generator.py new file mode 100644 index 0000000..71654a6 --- /dev/null +++ b/Solverz/code_printer/python/module/module_generator.py @@ -0,0 +1,230 @@ +from __future__ import annotations + +import inspect +import os +import re +from typing import Dict + +from Solverz.equation.equations import AE as SymAE, FDAE as SymFDAE, DAE as SymDAE +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 +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 + + +def render_modules(eqs: SymEquations, y, name, directory=None, numba=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'] + custom_func = dict() + custom_func.update(numerical_interface) + + def print_trigger_func_code(): + code_tfuc = dict() + trigger_func = parse_trigger_func(eqs.PARAM) + # rename trigger_func + for func_name, func in trigger_func.items(): + name0 = func.__name__ + func_code = inspect.getsource(func).replace(name0, func_name) + func_code = func_code.replace('np.', '') + code_tfuc[func_name] = func_code + return code_tfuc + + code_tfunc_dict = print_trigger_func_code() + + print('Complete!') + + eqn_parameter = {} + if isinstance(eqs, SymAE) and not isinstance(eqs, SymFDAE): + eqn_type = 'AE' + elif isinstance(eqs, SymFDAE): + eqn_type = 'FDAE' + eqn_parameter.update({'nstep': eqs.nstep}) + elif isinstance(eqs, SymDAE): + eqn_type = 'DAE' + eqn_parameter.update({'M': eqs.M}) + 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}) + print(f"Rendering python modules!") + render_as_modules(name, + code_dict, + eqn_type, + p, + eqn_parameter, + y, + [custom_func, 'numpy'], + numba, + directory) + print('Complete!') + + +def is_valid_python_module_name(module_name): + pattern = re.compile(r'^[a-zA-Z_][a-zA-Z0-9_]*$') + return bool(pattern.match(module_name)) + + +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 + os.makedirs(location, exist_ok=True) + + # Create an empty __init__.py file + init_path = os.path.join(location, "__init__.py") + with open(init_path, "w") as file: + file.write(initiate_code) + + # Create the file with the dependency code + module_path = os.path.join(location, "dependency.py") + with open(module_path, "w") as file: + file.write(dependency_code) + + # Create the file with the module code + module_path = os.path.join(location, "num_func.py") + with open(module_path, "w") as file: + file.write(module_code) + + save(auxiliary, f'{location}/param_and_setting.pkl') + + +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' + code += 'import time\n' + match eqn_type: + 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' + case 'FDAE': + try: + nstep = eqn_param['nstep'] + except KeyError as e: + raise ValueError("Cannot parse nstep attribute for FDAE object printing!") + 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' + 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' + case _: + raise ValueError(f'Unknown equation type {eqn_type}') + code += f'print("Compiling model {module_name}...")\n' + code += f'start = time.perf_counter()\n' + code += code_compile + code += f'end = time.perf_counter()\n' + code += "print(f'Compiling time elapsed: {end-start}s')\n" + return code + + +def print_module_code(code_dict: Dict[str, str], numba=False): + code = 'from .dependency import *\n' + code += """_data_ = setting["data"]\n""" + code += """_F_ = zeros_like(y, dtype=float64)""" + code += '\n\r\n' + for tfunc in code_dict['tfunc_dict'].values(): + if numba: + 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: + 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: + if numba: + code += '@njit(cache=True)\n' + code += sub_func + code += '\n\r\n' + return code + + +def print_dependency_code(modules): + code = "import os\n" + code += "current_module_dir = os.path.dirname(os.path.abspath(__file__))\n" + code += 'from Solverz import load\n' + code += 'auxiliary = load(f"{current_module_dir}\\\\param_and_setting.pkl")\n' + code += 'from numpy import *\n' + code += 'from numpy import abs\n' + code += 'from Solverz.num_api.custom_function import *\n' # import Solverz built-in func + code += 'from scipy.sparse import *\n' + code += 'from numba import njit\n' + code += 'setting = auxiliary["eqn_param"]\n' + code += 'row = setting["row"]\n' + code += 'col = setting["col"]\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): + if is_valid_python_module_name(name): + module_name = name + else: + raise ValueError("Invalid python module name!") + + init_code = print_init_code(eqn_type, module_name, eqn_param) + module_code = print_module_code(code_dict, numba=numba) + dependency_code = print_dependency_code(modules) + + create_python_module(module_name, + init_code, + module_code, + dependency_code, + {'p': p, 'eqn_param': eqn_param, 'vars': variables}, + directory) diff --git a/Solverz/code_printer/python/module/module_printer.py b/Solverz/code_printer/python/module/module_printer.py new file mode 100644 index 0000000..98e3e00 --- /dev/null +++ b/Solverz/code_printer/python/module/module_printer.py @@ -0,0 +1,131 @@ +from __future__ import annotations + +from Solverz.code_printer.python.utilities import * + + +# %% + + +def print_J(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"Jac matrix, with size ({eqn_size}*{var_addr.total_size}), not square") + fp = print_F_J_prototype(eqs_type, + 'J_', + 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)) + body.extend([Assignment(iVar('data', internal_use=True), + FunctionCall('inner_J', [symbols('_data_', real=True)] + var_list + param_list))]) + body.extend([Return(coo_2_csc(eqn_size, var_addr.total_size))]) + fd = FunctionDefinition.from_FunctionPrototype(fp, body) + return pycode(fd, fully_qualified_modules=False) + + +def print_inner_J(var_addr: Address, + PARAM: Dict[str, ParamBase], + jac: Jac, + 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_J', [symbols('_data_', real=True)] + args) + body = [] + + code_sub_inner_J_blocks = [] + count = 0 + addr_by_ele_0 = 0 + for eqn_name, jbs_row in jac.blocks.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_J0(t1, x) + body.append(Assignment(iVar('_data_', internal_use=True)[addr_by_ele], + FunctionCall(f'inner_J{int(count)}', SymbolsInDeri))) + + # def inner_J0(t1, x): + # return -t1 * pi * cos(pi * x) + 1 + fp1 = FunctionPrototype(real, f'inner_J{int(count)}', SymbolsInDeri) + body1 = [Return(rhs)] + fd1 = FunctionDefinition.from_FunctionPrototype(fp1, body1) + code_sub_inner_J_blocks.append(pycode(fd1, fully_qualified_modules=False)) + count += 1 + addr_by_ele_0 += jb.SpEleSize + temp = iVar('_data_', internal_use=True) + body.extend([Return(temp)]) + fd = FunctionDefinition.from_FunctionPrototype(fp, body) + return {'code_inner_J': pycode(fd, fully_qualified_modules=False), + 'code_sub_inner_J': code_sub_inner_J_blocks} + + +def print_F(eqs_type: str, + var_addr: Address, + PARAM: Dict[str, ParamBase], + nstep: int = 0): + fp = print_F_J_prototype(eqs_type, + 'F_', + 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)) + body.extend( + [Return(FunctionCall('inner_F', [symbols('_F_', real=True)] + var_list + param_list))]) + fd = FunctionDefinition.from_FunctionPrototype(fp, body) + return pycode(fd, fully_qualified_modules=False) + + +def print_inner_F(EQNs: Dict[str, Eqn], + EqnAddr: Address, + var_addr: Address, + PARAM: Dict[str, ParamBase], + 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_F', [symbols('_F_', real=True)] + args) + body = [] + body.extend(print_eqn_assignment(EQNs, + EqnAddr, + True)) + temp = iVar('_F_', internal_use=True) + body.extend([Return(temp)]) + fd = FunctionDefinition.from_FunctionPrototype(fp, body) + return pycode(fd, fully_qualified_modules=False) + + +def print_sub_inner_F(EQNs: Dict[str, Eqn]): + code_blocks = [] + count = 0 + for eqn_name in EQNs.keys(): + eqn = EQNs[eqn_name] + args = [] + for var in eqn.SYMBOLS.keys(): + args.append(symbols(var, real=True)) + fp = FunctionPrototype(real, f'inner_F{count}', args) + body = [Return(eqn.RHS)] + fd = FunctionDefinition.from_FunctionPrototype(fp, body) + count = count + 1 + code_blocks.append(pycode(fd, fully_qualified_modules=False)) + return code_blocks diff --git a/Solverz/code_printer/python/module/test/__init__.py b/Solverz/code_printer/python/module/test/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/Solverz/code_printer/python/module/test/test_module_generator.py b/Solverz/code_printer/python/module/test/test_module_generator.py new file mode 100644 index 0000000..4e79702 --- /dev/null +++ b/Solverz/code_printer/python/module/test/test_module_generator.py @@ -0,0 +1,281 @@ +import inspect +import os +import shutil +import sys + +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.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 +from Solverz.sym_algebra.symbols import idx, iVar, Para + +expected_dependency = r"""import os +current_module_dir = os.path.dirname(os.path.abspath(__file__)) +from Solverz import load +auxiliary = load(f"{current_module_dir}\\param_and_setting.pkl") +from numpy import * +from numpy import abs +from Solverz.num_api.custom_function import * +from scipy.sparse import * +from numba import njit +setting = auxiliary["eqn_param"] +row = setting["row"] +col = setting["col"] +p_ = auxiliary["p"] +y = auxiliary["vars"] +""" + +expected_init = r"""from .num_func import F_, J_ +from .dependency import setting, p_, y +import time +from Solverz.num_api.num_eqn import nAE +mdl = nAE(F_, J_, p_) +print("Compiling model test_eqn1...") +start = time.perf_counter() +mdl.F(y, p_) +mdl.J(y, p_) +end = time.perf_counter() +print(f'Compiling time elapsed: {end-start}s') +""" + + +def test_AE_module_printer(): + x = iVar('x', [1, 1]) + f1 = Eqn('f1', 2 * x[0] + x[1]) + f2 = Eqn('f2', x[0] ** 2 + sin(x[1])) + + F = AE([f1, f2]) + y = as_Vars([x]) + + current_file_path = os.path.abspath(__file__) + current_folder = os.path.dirname(current_file_path) + + test_folder_path = current_folder + '\\Solverz_testaabbccddeeffgghh' + + pyprinter = module_printer(F, + y, + 'test_eqn1', + directory=test_folder_path + '\\a_test_direc', + jit=True) + pyprinter.render() + + pyprinter1 = module_printer(F, + y, + 'test_eqn2', + directory=test_folder_path + '\\a_test_direc', + jit=False) + pyprinter1.render() + + sys.path.extend([test_folder_path + '\\a_test_direc']) + + from test_eqn1 import mdl, y + from test_eqn1.num_func import inner_F, inner_F0, inner_F1, inner_J + + F0 = mdl.F(y, mdl.p) + J0 = mdl.J(y, mdl.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) + + assert inspect.getsource( + inner_F) == '@njit(cache=True)\ndef inner_F(_F_, x):\n _F_[0:1] = inner_F0(x)\n _F_[1:2] = inner_F1(x)\n return _F_\n' + assert inspect.getsource(inner_F0.func_code) == '@njit(cache=True)\ndef inner_F0(x):\n return 2*x[0] + x[1]\n' + assert inspect.getsource( + inner_F1.func_code) == '@njit(cache=True)\ndef inner_F1(x):\n return x[0]**2 + sin(x[1])\n' + assert inspect.getsource( + inner_J) == '@njit(cache=True)\ndef inner_J(_data_, x):\n _data_[2:3] = inner_J0(x)\n _data_[3:4] = inner_J1(x)\n return _data_\n' + + # read dependency.py + with open(test_folder_path + '\\a_test_direc\\test_eqn1\\dependency.py', 'r', encoding='utf-8') as file: + file_content = file.read() + + assert file_content == expected_dependency + + with open(test_folder_path + '\\a_test_direc\\test_eqn1\\__init__.py', 'r', encoding='utf-8') as file: + file_content = file.read() + + assert file_content == expected_init + + from test_eqn2 import mdl, y + from test_eqn2.num_func import inner_F, inner_F0, inner_F1, inner_J + + F0 = mdl.F(y, mdl.p) + J0 = mdl.J(y, mdl.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) + + assert inspect.getsource( + inner_F) == 'def inner_F(_F_, x):\n _F_[0:1] = inner_F0(x)\n _F_[1:2] = inner_F1(x)\n return _F_\n' + assert inspect.getsource(inner_F0) == 'def inner_F0(x):\n return 2*x[0] + x[1]\n' + assert inspect.getsource(inner_F1) == 'def inner_F1(x):\n return x[0]**2 + sin(x[1])\n' + assert inspect.getsource( + inner_J) == 'def inner_J(_data_, x):\n _data_[2:3] = inner_J0(x)\n _data_[3:4] = inner_J1(x)\n return _data_\n' + + shutil.rmtree(test_folder_path) + + +expected_F = """def F_(t, y_, p_, y_0): + p = y_[0:82] + q = y_[82:164] + p_tag_0 = y_0[0:82] + q_tag_0 = y_0[82:164] + pb = p_["pb"].get_v_t(t) + qb = p_["qb"] + return inner_F(_F_, p, q, p_tag_0, q_tag_0, pb, qb) +""" + +expected_inner_F = """@njit(cache=True) +def inner_F(_F_, p, q, p_tag_0, q_tag_0, pb, qb): + _F_[0:81] = inner_F0(p, p_tag_0, q, q_tag_0) + _F_[81:162] = inner_F1(p, p_tag_0, q, q_tag_0) + _F_[162:163] = inner_F2(p, pb) + _F_[163:164] = inner_F3(q, qb) + return _F_ +""" + +expected_J = """def J_(t, y_, p_, y_0): + p = y_[0:82] + q = y_[82:164] + p_tag_0 = y_0[0:82] + q_tag_0 = y_0[82:164] + pb = p_["pb"].get_v_t(t) + qb = p_["qb"] + data = inner_J(_data_, p, q, p_tag_0, q_tag_0, pb, qb) + return coo_array((data, (row, col)), (164, 164)).tocsc() +""" + +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) + return _data_ +""" + + +def test_FDAE_module_generator(): + L = 51000 * 0.8 + p0 = 6621246.69079594 + q0 = 14 + + va = Integer(340) + D = 0.5901 + S = np.pi * (D / 2) ** 2 + lam = 0.03 + + dx = 500 + dt = 1.4706 + M = int(L / dx) + m1 = Model() + m1.p = Var('p', value=p0 * np.ones((M + 1,))) + m1.q = Var('q', value=q0 * np.ones((M + 1,))) + m1.p0 = AliasVar('p', init=m1.p) + m1.q0 = AliasVar('q', init=m1.q) + + m1.ae1 = Eqn('cha1', + m1.p[1:M + 1] - m1.p0[0:M] + va / S * (m1.q[1:M + 1] - m1.q0[0:M]) + + lam * va ** 2 * dx / (4 * D * S ** 2) * (m1.q[1:M + 1] + m1.q0[0:M]) * Abs( + m1.q[1:M + 1] + m1.q0[0:M]) / ( + m1.p[1:M + 1] + m1.p0[0:M])) + + m1.ae2 = Eqn('cha2', + m1.p0[1:M + 1] - m1.p[0:M] + va / S * (m1.q[0:M] - m1.q0[1:M + 1]) + + lam * va ** 2 * dx / (4 * D * S ** 2) * (m1.q[0:M] + m1.q0[1:M + 1]) * Abs( + m1.q[0:M] + m1.q0[1:M + 1]) / ( + m1.p[0:M] + m1.p0[1:M + 1])) + T = 5 * 3600 + pb1 = 1e6 + pb0 = 6621246.69079594 + pb_t = [pb0, pb0, pb1, pb1] + tseries = [0, 1000, 1000 + 10 * dt, T] + m1.pb = TimeSeriesParam('pb', + v_series=pb_t, + time_series=tseries) + m1.qb = Param('qb', q0) + m1.bd1 = Eqn('bd1', m1.p[0] - m1.pb) + m1.bd2 = Eqn('bd2', m1.q[M] - m1.qb) + fdae, y0 = m1.create_instance() + + current_file_path = os.path.abspath(__file__) + current_folder = os.path.dirname(current_file_path) + + test_folder_path = current_folder + '\\Solverz_testaabbccddeeffgghh' + + pyprinter = module_printer(fdae, + y0, + 'test_fdae', + directory=test_folder_path, + jit=True) + pyprinter.render() + + sys.path.extend([test_folder_path]) + + from test_fdae.num_func import F_, J_, inner_F, inner_J + + assert inspect.getsource(F_) == expected_F + assert inspect.getsource(J_) == expected_J + assert inspect.getsource(inner_F) == expected_inner_F + assert inspect.getsource(inner_J) == expected_inner_J + + shutil.rmtree(test_folder_path) + + +expected_F1 = """def F_(t, y_, p_): + h = y_[0:1] + v = y_[1:2] + return inner_F(_F_, h, v) +""" + +expected_inner_F1 = """@njit(cache=True) +def inner_F(_F_, h, v): + _F_[0:1] = inner_F0(v) + _F_[1:2] = inner_F1() + return _F_ +""" + +expected_J1 = """def J_(t, y_, p_): + h = y_[0:1] + v = y_[1:2] + data = inner_J(_data_, h, v) + return coo_array((data, (row, col)), (2, 2)).tocsc() +""" + +expected_inner_J1 = """@njit(cache=True) +def inner_J(_data_, h, v): + return _data_ +""" + + +def test_DAE_module_generator(): + m = Model() + m.h = Var('h', 0) + m.v = Var('v', 20) + m.f1 = Ode('f1', f=m.v, diff_var=m.h) + m.f2 = Ode('f2', f=-9.8, diff_var=m.v) + bball, 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_testaabbccddeeffgghh' + + pyprinter = module_printer(bball, + y0, + 'test_dae', + directory=test_folder_path, + jit=True) + pyprinter.render() + + sys.path.extend([test_folder_path]) + + from test_dae.num_func import F_, J_, inner_F, inner_J + + assert inspect.getsource(F_) == expected_F1 + assert inspect.getsource(J_) == expected_J1 + assert inspect.getsource(inner_F) == expected_inner_F1 + assert inspect.getsource(inner_J) == expected_inner_J1 + + shutil.rmtree(test_folder_path) diff --git a/Solverz/code_printer/python/module/test/test_module_printer.py b/Solverz/code_printer/python/module/test/test_module_printer.py new file mode 100644 index 0000000..d6c31b0 --- /dev/null +++ b/Solverz/code_printer/python/module/test/test_module_printer.py @@ -0,0 +1,284 @@ +import pytest +import numpy as np +import re + +from sympy import symbols, pycode, Integer +from sympy.codegen.ast import FunctionCall as SpFuncCall, Assignment + +from Solverz.sym_algebra.symbols import iVar, Para +from Solverz.equation.eqn import Eqn, Ode +from Solverz.equation.param import Param, TimeSeriesParam +from Solverz.equation.jac import Jac, JacBlock +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) + +expected = """def J_(t, y_, p_): + omega = y_[0:10] + delta = y_[10:15] + x = y_[15:18] + y = y_[18:25] + ax = p_["ax"] + lam = p_["lam"] + G6 = p_["G6"].get_v_t(t) + ax = ax_trigger_func(x) + data = inner_J(_data_, omega, delta, x, y, ax, lam, G6) + return coo_array((data, (row, col)), (25, 25)).tocsc() +""".strip() +expected1 = """def J_(t, y_, p_, 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"] + G6 = p_["G6"].get_v_t(t) + ax = ax_trigger_func(x) + data = inner_J(_data_, omega, delta, x, y, omega_tag_0, delta_tag_0, x_tag_0, y_tag_0, ax, lam, G6) + return coo_array((data, (row, col)), (25, 25)).tocsc() +""".strip() + + +def test_print_J(): + eqs_type = 'DAE' + 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 + G6 = TimeSeriesParam('G6', + [0, 1, 2, 3], + [0, 100, 200, 300]) + Pdict["G6"] = G6 + + with pytest.raises(ValueError, match=re.escape("Jac matrix, with size (20*25), not square")): + print_J(eqs_type, + 20, + VarAddr, + Pdict) + + assert print_J(eqs_type, + 25, + VarAddr, + Pdict) == expected + + assert print_J("FDAE", + 25, + VarAddr, + Pdict, + 1) == expected1 + + +expected2 = """def inner_J(_data_, omega, x, y, z, ax, lam, G6): + _data_[0:3] = inner_J0(y) + _data_[3:12] = inner_J1(y) + return _data_ +""".strip() +expected3 = """def inner_J0(y): + return y*ones(3) +""".strip() +expected4 = """def inner_J1(y): + return y**2 +""".strip() +expected5 = """def inner_J(_data_, omega, x, y, z, omega_tag_0, x_tag_0, y_tag_0, z_tag_0, ax, lam, G6): + _data_[0:3] = inner_J0(y) + _data_[3:12] = inner_J1(y) + return _data_ +""".strip() + + +def test_print_inner_J(): + jac = Jac() + x = iVar('x') + y = iVar('y') + z = iVar('z') + omega = iVar('omega') + jac.add_block('a', + x[0], + JacBlock('a', + slice(0, 3), + x[0], + np.array([1]), + slice(1, 2), + iVar('y'), + np.array([1]))) + jac.add_block('b', + omega[1:4], + JacBlock('b', + slice(3, 12), + omega[4:13], + np.ones(9), + slice(3, 100), + iVar('y') ** 2, + np.ones(9))) + jac.add_block('b', + y, + JacBlock('b', + slice(3, 12), + z, + np.ones(9), + slice(100, 109), + Integer(1), + np.ones(1))) + + VarAddr = Address() + VarAddr.add('omega', 97) + VarAddr.add('x', 1) + VarAddr.add('y', 1) + VarAddr.add('z', 9) + 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 + G6 = TimeSeriesParam('G6', + [0, 1, 2, 3], + [0, 100, 200, 300]) + Pdict["G6"] = G6 + + code_dict = print_inner_J(VarAddr, + Pdict, + jac) + + assert code_dict['code_inner_J'] == expected2 + assert code_dict['code_sub_inner_J'][0] == expected3 + assert code_dict['code_sub_inner_J'][1] == expected4 + + code_dict = print_inner_J(VarAddr, + Pdict, + jac, + 1) + + assert code_dict['code_inner_J'] == expected5 + + +expected6 = """def F_(t, y_, p_): + omega = y_[0:10] + delta = y_[10:15] + x = y_[15:18] + y = y_[18:25] + ax = p_["ax"] + lam = p_["lam"] + G6 = p_["G6"].get_v_t(t) + ax = ax_trigger_func(x) + return inner_F(_F_, omega, delta, x, y, ax, lam, G6) +""".strip() + + +def test_print_F(): + eqs_type = 'DAE' + 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 + G6 = TimeSeriesParam('G6', + [0, 1, 2, 3], + [0, 100, 200, 300]) + Pdict["G6"] = G6 + assert print_F(eqs_type, + VarAddr, + Pdict) == expected6 + + +expected7 = """def inner_F(_F_, omega, delta, x, y, ax, lam, G6): + _F_[0:10] = inner_F0(delta, omega) + _F_[10:15] = inner_F1(lam, x, y) + return _F_ +""".strip() +expected8 = """def inner_F(_F_, omega, delta, x, y, omega_tag_0, delta_tag_0, x_tag_0, y_tag_0, ax, lam, G6): + _F_[0:10] = inner_F0(delta, omega) + _F_[10:15] = inner_F1(lam, x, y) + return _F_ +""".strip() + + +def test_print_inner_F(): + EQNs = dict() + omega = iVar('omega') + delta = iVar('delta') + x = iVar('x') + y = iVar('y') + lam = Para('lam') + EQNs['a'] = Eqn('a', omega * delta) + EQNs['b'] = Ode('b', x + y * lam, diff_var=x) + EqnAddr = Address() + EqnAddr.add('a', 10) + EqnAddr.add('b', 5) + 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 + G6 = TimeSeriesParam('G6', + [0, 1, 2, 3], + [0, 100, 200, 300]) + Pdict["G6"] = G6 + assert print_inner_F(EQNs, + EqnAddr, + VarAddr, + Pdict) == expected7 + + assert print_inner_F(EQNs, + EqnAddr, + VarAddr, + Pdict, + 1) == expected8 + + +expected9 = """def inner_F0(delta, omega): + return delta*omega +""".strip() +expected10 = """def inner_F1(lam, x, y): + return lam*y + x +""".strip() + + +def test_print_sub_inner_F(): + EQNs = dict() + omega = iVar('omega') + delta = iVar('delta') + x = iVar('x') + y = iVar('y') + lam = Para('lam') + EQNs['a'] = Eqn('a', omega * delta) + EQNs['b'] = Ode('b', x + y * lam, diff_var=x) + + assert print_sub_inner_F(EQNs)[0] == expected9 + assert print_sub_inner_F(EQNs)[1] == expected10 diff --git a/Solverz/code_printer/python/tests/__init__.py b/Solverz/code_printer/python/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/Solverz/code_printer/python/tests/test_utilities.py b/Solverz/code_printer/python/tests/test_utilities.py new file mode 100644 index 0000000..bd0e6ab --- /dev/null +++ b/Solverz/code_printer/python/tests/test_utilities.py @@ -0,0 +1,211 @@ +import pytest +import numpy as np + +from sympy import symbols, pycode +from sympy.codegen.ast import FunctionCall as SpFuncCall, Assignment + +from Solverz.code_printer.python.utilities import FunctionCall, parse_p, parse_trigger_func, \ + print_var, print_param, print_trigger, print_F_J_prototype, print_eqn_assignment, zeros +from Solverz.sym_algebra.symbols import iVar, Para +from Solverz.equation.eqn import Eqn, Ode +from Solverz.equation.param import Param, TimeSeriesParam +from Solverz.utilities.address import Address +from sympy.codegen.ast import FunctionDefinition, Return + +expected = """ + The `args` parameter passed to sympy.codegen.ast.FunctionCall should not contain str, which may cause sympy parsing error. + For example, the sympy.codegen.ast.FunctionCall parses str E in args to math.e! Please use sympy.Symbol objects instead. + """ + + +def test_FunctionCall(): + E = symbols('E') + assert pycode(FunctionCall('a', [E])) == 'a(E)' + assert pycode(SpFuncCall('a', ['E'])) == 'a(math.e)' + + with pytest.raises(ValueError, match=expected): + assert pycode(FunctionCall('a', ['E'])) == 'a(E)' + + +def test_parse_p(): + Pdict = dict() + Pdict["lam"] = Param('lam', [1, 1, 1]) + Pdict["G6"] = TimeSeriesParam('G6', + [0, 1, 2, 3], + [0, 100, 200, 300]) + p = parse_p(Pdict) + assert "lam" in p + np.testing.assert_allclose(p["lam"], np.array([1, 1, 1])) + assert "G6" in p + np.testing.assert_allclose(p["G6"].value, 0) + np.testing.assert_allclose(p["G6"].get_v_t(40), 0.4) + + +def test_parse_trigger_func(): + ax = Param('ax', + [0, 1], + triggerable=True, + trigger_var=['x'], + trigger_fun=np.sin) + lam = Param('lam', [1, 1, 1]) + G6 = TimeSeriesParam('G6', + [0, 1, 2, 3], + [0, 100, 200, 300]) + func_dict = parse_trigger_func({'ax': ax, 'lam': lam, 'G6': G6}) + assert "ax_trigger_func" in func_dict + np.testing.assert_allclose(np.sin(1), func_dict["ax_trigger_func"](1)) + + +def test_print_var(): + va = Address() + va.add('delta', 10) + va.add('omega', 3) + va.add('p', 2) + var_dec, var_list = print_var(va, 0) + y_ = iVar('y_', internal_use=True) + delta = iVar('delta') + omega = iVar('omega') + p = iVar('p') + assert var_dec == [Assignment(delta, y_[0:10]), + Assignment(omega, y_[10:13]), + Assignment(p, y_[13:15])] + assert var_list == [iVar('delta'), + iVar('omega'), + iVar('p')] + + var_dec, var_list = print_var(va, 1) + suffix = '0' + delta_tag_0 = iVar('delta' + '_tag_' + suffix) + omega_tag_0 = iVar('omega' + '_tag_' + suffix) + p_tag_0 = iVar('p' + '_tag_' + suffix) + y_0 = iVar('y_' + suffix, internal_use=True) + assert var_dec == [Assignment(delta, y_[0:10]), + Assignment(omega, y_[10:13]), + Assignment(p, y_[13:15]), + Assignment(delta_tag_0, y_0[0:10]), + Assignment(omega_tag_0, y_0[10:13]), + Assignment(p_tag_0, y_0[13:15])] + assert var_list == [delta, + omega, + p, + delta_tag_0, + omega_tag_0, + p_tag_0] + + +def test_print_param(): + Pdict = dict() + lam = Param('lam', [1, 1, 1]) + Pdict["lam"] = lam + G6 = TimeSeriesParam('G6', + [0, 1, 2, 3], + [0, 100, 200, 300]) + Pdict["G6"] = G6 + Area = Param('Area', + [1, 1, 1], + is_alias=True) + Pdict["Area"] = Area + param_declaration, param_list = print_param(Pdict) + assert pycode(param_declaration[0]) == 'lam = p_["lam"]' + assert pycode(param_declaration[1]) == 'G6 = p_["G6"].get_v_t(t)' + assert param_list[0] == Para('lam') + assert param_list[1] == Para('G6') + + +def test_print_trigger(): + 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 + G6 = TimeSeriesParam('G6', + [0, 1, 2, 3], + [0, 100, 200, 300]) + Pdict["G6"] = G6 + tri_decla = print_trigger(Pdict) + assert pycode(tri_decla[0]) == 'ax = ax_trigger_func(x)' + + +expected1 = """def F_(y_, p_): + return a +""".strip() +expected2 = """def J_(y_, p_): + return a +""".strip() +expected3 = """def F_(t, y_, p_): + return a +""".strip() +expected4 = """def J_(t, y_, p_): + return a +""".strip() +expected5 = """def J_(t, y_, p_, y_0): + return a +""".strip() +expected6 = """def F_(t, y_, p_, y_0): + return a +""".strip() + + +def test_print_F_J_prototype(): + body = [] + body.extend([Return(iVar('a'))]) + + fp = print_F_J_prototype('AE', 'F_', 0) + fd = FunctionDefinition.from_FunctionPrototype(fp, body) + assert pycode(fd) == expected1 + + fp = print_F_J_prototype('AE', 'J_', 0) + fd = FunctionDefinition.from_FunctionPrototype(fp, body) + assert pycode(fd) == expected2 + + fp = print_F_J_prototype('DAE', 'F_', 0) + fd = FunctionDefinition.from_FunctionPrototype(fp, body) + assert pycode(fd) == expected3 + + fp = print_F_J_prototype('DAE', 'J_', 0) + fd = FunctionDefinition.from_FunctionPrototype(fp, body) + assert pycode(fd) == expected4 + + with pytest.raises(TypeError, match="Only FDAE supports the case where nstep > 0! Not DAE"): + fp = print_F_J_prototype('DAE', 'J_', 1) + + with pytest.raises(ValueError, match="Func name A not supported!"): + fp = print_F_J_prototype('DAE', 'A', 1) + + with pytest.raises(ValueError, match="nstep of FDAE should be greater than 0!"): + fp = print_F_J_prototype('FDAE', 'J_', 0) + + fp = print_F_J_prototype('FDAE', 'J_', 1) + fd = FunctionDefinition.from_FunctionPrototype(fp, body) + assert pycode(fd) == expected5 + + fp = print_F_J_prototype('FDAE', 'F_', 1) + fd = FunctionDefinition.from_FunctionPrototype(fp, body) + assert pycode(fd) == expected6 + + +def test_eqn_assignment(): + _F_ = iVar('_F_', internal_use=True) + EQNs = dict() + omega = iVar('omega') + delta = iVar('delta') + x = iVar('x') + y = iVar('y') + lam = Para('lam') + EQNs['a'] = Eqn('a', omega * delta) + EQNs['b'] = Ode('b', x + y * lam, diff_var=x) + EqnAddr = Address() + EqnAddr.add('a', 10) + EqnAddr.add('b', 5) + eqn_decl = print_eqn_assignment(EQNs, EqnAddr) + assert eqn_decl[0] == Assignment(_F_, zeros(15)) + assert eqn_decl[1] == Assignment(_F_[0:10], delta*omega) + assert eqn_decl[2] == Assignment(_F_[10:15], lam*y+x) + + eqn_decl = print_eqn_assignment(EQNs, EqnAddr, True) + assert eqn_decl[0] == Assignment(_F_[0:10], FunctionCall('inner_F0', [delta, omega])) + assert eqn_decl[1] == Assignment(_F_[10:15], FunctionCall('inner_F1', [lam, x, y])) diff --git a/Solverz/code_printer/python/utilities.py b/Solverz/code_printer/python/utilities.py new file mode 100644 index 0000000..7d0649c --- /dev/null +++ b/Solverz/code_printer/python/utilities.py @@ -0,0 +1,218 @@ +from __future__ import annotations + +import warnings +from typing import Union, List, Dict, Callable +import builtins +from datetime import datetime + +import numpy as np +from sympy.codegen.ast import Assignment, AddAugmentedAssignment +from sympy import pycode, symbols, Function, Symbol, Expr, Number as SymNumber +from sympy.codegen.ast import real, FunctionPrototype, FunctionDefinition, Return, FunctionCall as SymFuncCall +from sympy.utilities.lambdify import _import, _module_present, _get_namespace +from scipy.sparse import sparray +from numbers import Number + +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.param import TimeSeriesParam, ParamBase +from Solverz.sym_algebra.symbols import iVar, SolDict, Para, idx, IdxSymBasic +from Solverz.sym_algebra.functions import Arange +from Solverz.utilities.address import Address +from Solverz.utilities.type_checker import is_number +from Solverz.num_api.custom_function import numerical_interface +from Solverz.num_api.num_eqn import nAE, nFDAE, nDAE + + +# %% + +def FunctionCall(name, args): + if not isinstance(args, list): + raise TypeError("args should be a list.") + for i in range(len(args)): + if isinstance(args[i], str): + raise ValueError(""" + The `args` parameter passed to sympy.codegen.ast.FunctionCall should not contain str, which may cause sympy parsing error. + For example, the sympy.codegen.ast.FunctionCall parses str E in args to math.e! Please use sympy.Symbol objects instead. + """) + return SymFuncCall(name, args) + + +def parse_p(PARAM: Dict[str, ParamBase]): + p = dict() + for param_name, param in PARAM.items(): + if isinstance(param, TimeSeriesParam): + p.update({param_name: param}) + else: + p.update({param_name: param.v}) + return p + + +def parse_trigger_func(PARAM: Dict[str, ParamBase]) -> Dict[str, Callable]: + func = dict() + for para_name, param in PARAM.items(): + if param.trigger_fun is not None: + func.update({para_name + '_trigger_func': param.trigger_fun}) + + return func + + +def _print_var_parser(var_name, suffix: str, var_addr_slice): + y = iVar('y_' + suffix, internal_use=True) + if suffix == '': + return Assignment(iVar(var_name + suffix), y[var_addr_slice]) + else: + return Assignment(iVar(var_name + '_tag_' + suffix), y[var_addr_slice]) + pass + + +def print_var(var_addr: Address, nstep): + var_declaration = [] + suffix = [''] + [f'{i}' for i in range(nstep)] + var_list = [] + for i in suffix: + for var_name in var_addr.v.keys(): + var_addr_slice = var_addr[var_name] + var_assign = _print_var_parser(var_name, i, var_addr_slice) + var_declaration.append(var_assign) + var_list.append(var_assign.lhs) + return var_declaration, var_list + + +def print_param(PARAM: Dict[str, ParamBase]): + param_declaration = [] + p = SolDict('p_') + param_list = [] + for param_name, param in PARAM.items(): + if not param.is_alias: + if isinstance(param, TimeSeriesParam): + param_assign = Assignment(Para(param_name), + FunctionCall(f'p_["{param_name}"].get_v_t', [Symbol('t')])) + else: + param_assign = Assignment(Para(param_name), + p[param_name]) + param_declaration.append(param_assign) + param_list.append(param_assign.lhs) + + return param_declaration, param_list + + +def _print_trigger_func(para_name, trigger_var: List[str]): + return Assignment(Para(para_name), + FunctionCall(para_name + '_trigger_func', [symbols(var) for var in trigger_var])) + + +def print_trigger(PARAM: Dict[str, ParamBase]): + trigger_declaration = [] + for name, param in PARAM.items(): + if param.triggerable: + trigger_declaration.append(_print_trigger_func(name, + param.trigger_var)) + return trigger_declaration + + +def print_F_J_prototype(eqs_type: str, func_name: str, nstep=0): + if func_name not in ['F_', 'J_']: + raise ValueError(f"Func name {func_name} not supported!") + t, y_, p_ = symbols('t y_ p_', real=True) + if eqs_type == 'DAE' or eqs_type == 'FDAE': + args = [t, y_, p_] + elif eqs_type == 'AE': + args = [y_, p_] + 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 + + fp = FunctionPrototype(real, func_name, args + xtra_args) + return fp + + +def print_eqn_assignment(EQNs: Dict[str, Eqn], + EqnAddr: Address, + module_printer=False): + _F_ = iVar('_F_', internal_use=True) + eqn_declaration = [] + if not module_printer: + eqn_declaration.append(Assignment(_F_, zeros(EqnAddr.total_size, ))) + # Whereas in module printer, the _F_ is predeclared for only once. + count = 0 + for eqn_name in EQNs.keys(): + eqn_address = EqnAddr[eqn_name] + eqn = EQNs[eqn_name] + if module_printer: + 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}") + _F_ = iVar('_F_', internal_use=True) + eqn_declaration.append(Assignment(_F_[eqn_address], + FunctionCall(f'inner_F{int(count)}', list(eqn.SYMBOLS.values())))) + + count = count + 1 + else: + eqn_declaration.append(Assignment(_F_[eqn_address], eqn.RHS)) + return eqn_declaration + + +class coo_2_csc(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, (row, col)), ({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.") + + def _numpycode(self, printer, **kwargs): + return f'coo_array({printer._print(self.args[0])})' + + def _pythoncode(self, printer, **kwargs): + return self._numpycode(printer, **kwargs) + + +class extend(Function): + + def _numpycode(self, printer, **kwargs): + return f'{printer._print(self.args[0])}.extend({printer._print(self.args[1])})' + + def _pythoncode(self, printer, **kwargs): + return self._numpycode(printer, **kwargs) + + +class zeros(Function): + # print zeros(6,6) as zeros((6,6)) + # or zeros(6,) as zeros((6,)) + def _numpycode(self, printer, **kwargs): + if len(self.args) == 2: + temp1 = printer._print(self.args[0]) + temp2 = printer._print(self.args[1]) + return r'zeros((' + temp1 + ', ' + temp2 + r'))' + elif len(self.args) == 1: + temp = printer._print(self.args[0]) + return r'zeros((' + temp + ', ' + r'))' + + def _pythoncode(self, printer, **kwargs): + return self._numpycode(printer, **kwargs) diff --git a/Solverz/code_printer/test/test_py_printer.py b/Solverz/code_printer/test/test_py_printer.py deleted file mode 100644 index 68f4b3c..0000000 --- a/Solverz/code_printer/test/test_py_printer.py +++ /dev/null @@ -1,266 +0,0 @@ -import inspect -import os -import shutil -import sys - -import numpy as np -from sympy import symbols, pycode, Integer - -from Solverz import as_Vars, Eqn, AE, sin, made_numerical -from Solverz.code_printer.make_module import module_printer -from Solverz.code_printer.py_printer import _parse_jac_eqn_address, _parse_jac_var_address, _parse_jac_data, \ - print_J_block, _print_var_parser -from Solverz.sym_algebra.symbols import idx, iVar, Para - - -def test_address_parser(): - eqn_address = slice(0, 1) - assert pycode(_parse_jac_eqn_address(eqn_address, 1, False)) == '0' - assert _parse_jac_eqn_address(eqn_address, 2, False).__repr__() == 'slice(0, 1, None)' - assert pycode(_parse_jac_eqn_address(eqn_address, 1, True)) == '[0]' - assert pycode(_parse_jac_eqn_address(eqn_address, 2, True)) == 'arange(0, 1)[value_coo.row]' - - eqn_address = slice(0, 5) - assert pycode(_parse_jac_eqn_address(eqn_address, 1, False)) == 'arange(0, 5)' - assert pycode(_parse_jac_eqn_address(eqn_address, 1, True)) == 'arange(0, 5)' - assert _parse_jac_eqn_address(eqn_address, 2, False).__repr__() == 'slice(0, 5, None)' - assert _parse_jac_eqn_address(eqn_address, 2, True).__repr__() == 'arange(0, 5)[value_coo.row]' - - v_address = slice(0, 1) - assert pycode(_parse_jac_var_address(v_address, 1, None, False)) == 'arange(0, 1)' - assert pycode(_parse_jac_var_address(v_address, 1, None, True)) == 'arange(0, 1)' - assert _parse_jac_var_address(v_address, 2, None, False).__repr__() == 'slice(0, 1, None)' - assert pycode(_parse_jac_var_address(v_address, 2, None, True)) == 'arange(0, 1)[value_coo.col]' - - v_address = slice(0, 3) - # _parse_jac_var_address(v_address, 2, v_idx, True) - assert _parse_jac_var_address(v_address, 2, None, False).__repr__() == 'slice(0, 3, None)' - assert _parse_jac_var_address(v_address, 2, None, True).__repr__() == 'arange(0, 3)[value_coo.col]' - assert pycode(_parse_jac_var_address(v_address, 1, None, True)) == 'arange(0, 3)' - assert pycode(_parse_jac_var_address(v_address, 1, None, False)) == 'arange(0, 3)' - assert pycode(_parse_jac_var_address(v_address, 0, 2, True, 3)) == '3*[2]' - assert pycode(_parse_jac_var_address(v_address, 1, 2, True, 3)) == '3*[2]' - - M = idx('M') - assert _parse_jac_var_address(v_address, 2, M, True).__repr__() == 'arange(0, 3)[M[value_coo.col]]' - assert pycode(_parse_jac_var_address(v_address, 2, M, False)) == 'arange(0, 3)[M]' - assert _parse_jac_var_address(v_address, 1, M, True).__repr__() == 'arange(0, 3)[M]' - assert pycode(_parse_jac_var_address(v_address, 1, M, False)) == 'arange(0, 3)[M]' - - v_idx = [0, 1, 2] - assert pycode(_parse_jac_var_address(v_address, 1, v_idx, True)) == 'arange(0, 3)' - assert pycode(_parse_jac_var_address(v_address, 1, v_idx, False)) == 'arange(0, 3)' - assert pycode(_parse_jac_var_address(v_address, 2, v_idx, True)) == 'arange(0, 3)[value_coo.col]' - assert _parse_jac_var_address(v_address, 2, v_idx, False).__repr__() == 'slice(0, 3, None)' - - v_address = slice(0, 10) - v_idx = [0, 2, 4, 6] - assert pycode(_parse_jac_var_address(v_address, 1, v_idx, True)) == '[0, 2, 4, 6]' - assert pycode(_parse_jac_var_address(v_address, 1, v_idx, False)) == '[0, 2, 4, 6]' - assert pycode(_parse_jac_var_address(v_address, 2, v_idx, True)) == 'array([0,2,4,6])[value_coo.col]' - assert pycode(_parse_jac_var_address(v_address, 2, v_idx, False)) == '[0, 2, 4, 6]' - - v_idx = 2 - assert pycode(_parse_jac_var_address(v_address, 1, v_idx, True)) == '[2]' - assert pycode(_parse_jac_var_address(v_address, 1, v_idx, False)) == '2' - assert _parse_jac_var_address(v_address, 2, v_idx, True).__repr__() == 'array([2])[value_coo.col]' - assert pycode(_parse_jac_var_address(v_address, 2, v_idx, False)) == '2' - - x, y = symbols('x, y') - assert pycode(_parse_jac_data(5, 1, x * y)) == 'x*y' - assert pycode(_parse_jac_data(5, 0, x * y)) == '5*((x*y).tolist())' - assert pycode(_parse_jac_data(5, 0, 3, rhs_v_type='Number')) == '5*[3]' - assert pycode(_parse_jac_data(5, 2, 3)) == 'value_coo.data' - - -def test_var_parser(): - assert pycode(_print_var_parser('x', '0', slice(0, 1))) == 'x_tag_0 = y_0[0:1]' - assert pycode(_print_var_parser('x', '0', slice(3, 6))) == 'x_tag_0 = y_0[3:6]' - - -# def test_F_printer(): -# F = _print_F_assignment(slice(20, 21), iVar('x')) -# assert pycode(F) == '_F_[20:21] = x' -# F = _print_F_assignment(slice(20, 23), iVar('x')) -# assert pycode(F) == '_F_[20:23] = x' - - -def test_J_block_printer(): - # dense - Jb = print_J_block(slice(0, 3), - slice(3, 6), - 0, - None, - iVar('omega_b'), - False) - assert pycode(Jb) == '[J_[arange(0, 3),arange(3, 6)] += omega_b]' - Jb = print_J_block(slice(3, 6), - slice(12, 21), - 1, - idx('g'), - -iVar('Ixg') / iVar('Tj'), - False) - assert pycode(Jb) == '[J_[arange(3, 6),arange(12, 21)[g]] += -Ixg/Tj]' - Jb = print_J_block(slice(12, 15), - slice(6, 9), - 0, - None, - 1, - False) - assert pycode(Jb) == '[J_[arange(12, 15),arange(6, 9)] += 1]' - Jb = print_J_block(slice(18, 24), - slice(12, 21), - 2, - None, - iVar('G')[idx('ng'), :], - False) - assert pycode(Jb) == '[J_[18:24,12:21] += G[ng,:]]' - - Jb = print_J_block(slice(30, 31), - slice(0, 10), - 2, - None, - Para('G', dim=2), - False) - assert pycode(Jb) == '[J_[30:31,0:10] += G]' - - # sparse - Jb = print_J_block(slice(39, 47), - slice(15, 26), - 0, - [0, 1, 2, 3, 4, 5, 9, 10], - Integer(1), - True, - rhs_v_dtpe='Number') - assert pycode( - Jb) == '[row.extend(arange(39, 47)), col.extend([15, 16, 17, 18, 19, 20, 24, 25]), data.extend(8*[1])]' - - Jb = print_J_block(slice(248, 259), - slice(269, 280), - 0, - None, - Para('dt') * Para('va') ** 2 / Para('S'), - True) - assert pycode( - Jb) == '[row.extend(arange(248, 259)), col.extend(arange(269, 280)), data.extend(11*((dt*va**2/S).tolist()))]' - - Jb = print_J_block(slice(281, 292), - slice(281, 292), - 1, - None, - Para('dx') + iVar('p'), - True) - assert pycode(Jb) == '[row.extend(arange(281, 292)), col.extend(arange(281, 292)), data.extend(dx + p)]' - - Jb = print_J_block(slice(0, 3), - slice(3, 6), - 2, - None, - iVar('omega_b'), - True) - assert pycode( - Jb) == '[value_coo = coo_array(omega_b), row.extend(arange(0, 3)[value_coo.row]), col.extend(arange(3, 6)[value_coo.col]), data.extend(value_coo.data)]' - - Jb = print_J_block(slice(3, 6), - slice(12, 21), - 2, - idx('g'), - Para('G'), - True) - assert pycode( - Jb) == '[value_coo = coo_array(G), row.extend(arange(3, 6)[value_coo.row]), col.extend(arange(12, 21)[g[value_coo.col]]), data.extend(value_coo.data)]' - - Jb = print_J_block(slice(39, 47), - slice(15, 26), - 2, - [0, 1, 2, 3, 4, 5, 9, 10], - Para('G'), - True) - assert pycode( - Jb) == '[value_coo = coo_array(G), row.extend(arange(39, 47)[value_coo.row]), col.extend(array([15,16,17,18,19,20,24,25])[value_coo.col]), data.extend(value_coo.data)]' - - Jb = print_J_block(slice(39, 47), - slice(15, 26), - 2, - slice(0, 5), - Para('G'), - True) - assert pycode( - Jb) == '[value_coo = coo_array(G), row.extend(arange(39, 47)[value_coo.row]), col.extend(arange(15, 20)[value_coo.col]), data.extend(value_coo.data)]' - - -def test_py_printer(): - x = iVar('x', [1, 1]) - f1 = Eqn('f1', 2 * x[0] + x[1]) - f2 = Eqn('f2', x[0] ** 2 + sin(x[1])) - - F = AE([f1, f2]) - y = as_Vars([x]) - - current_file_path = os.path.abspath(__file__) - current_folder = os.path.dirname(current_file_path) - - test_folder_path = current_folder + '\\Solverz_testaabbccddeeffgghh' - - pyprinter = module_printer(F, - y, - 'test_eqn1', - directory=test_folder_path + '\\a_test_direc', - jit=True) - pyprinter.render() - - pyprinter1 = module_printer(F, - y, - 'test_eqn2', - directory=test_folder_path + '\\a_test_direc', - jit=False) - pyprinter1.render() - - sys.path.extend([test_folder_path + '\\a_test_direc']) - - from test_eqn1 import mdl, y - from test_eqn1.num_func import inner_F, inner_F0, inner_F1, inner_J - - F0 = mdl.F(y, mdl.p) - J0 = mdl.J(y, mdl.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) - - assert inspect.getsource( - inner_F) == '@njit(cache=True)\ndef inner_F(_F_, x):\n _F_[0:1] = inner_F0(x)\n _F_[1:2] = inner_F1(x)\n return _F_\n' - assert inspect.getsource(inner_F0.func_code) == '@njit(cache=True)\ndef inner_F0(x):\n return 2*x[0] + x[1]\n' - assert inspect.getsource( - inner_F1.func_code) == '@njit(cache=True)\ndef inner_F1(x):\n return x[0]**2 + sin(x[1])\n' - assert inspect.getsource( - inner_J) == '@njit(cache=True)\ndef inner_J(_data_, x):\n _data_[2:3] = inner_J0(x)\n _data_[3:4] = inner_J1(x)\n return _data_\n' - - from test_eqn2 import mdl, y - from test_eqn2.num_func import inner_F, inner_F0, inner_F1, inner_J - - F0 = mdl.F(y, mdl.p) - J0 = mdl.J(y, mdl.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) - - assert inspect.getsource( - inner_F) == 'def inner_F(_F_, x):\n _F_[0:1] = inner_F0(x)\n _F_[1:2] = inner_F1(x)\n return _F_\n' - assert inspect.getsource(inner_F0) == 'def inner_F0(x):\n return 2*x[0] + x[1]\n' - assert inspect.getsource(inner_F1) == 'def inner_F1(x):\n return x[0]**2 + sin(x[1])\n' - assert inspect.getsource( - inner_J) == 'def inner_J(_data_, x):\n _data_[2:3] = inner_J0(x)\n _data_[3:4] = inner_J1(x)\n return _data_\n' - - shutil.rmtree(test_folder_path) - - -def test_made_numerical(): - x = iVar('x', [1, 1]) - f1 = Eqn('f1', 2 * x[0] + x[1]) - f2 = Eqn('f2', x[0] ** 2 + sin(x[1])) - - F = AE([f1, f2]) - y = as_Vars([x]) - 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) diff --git a/Solverz/equation/eqn.py b/Solverz/equation/eqn.py index 8676de8..b450c6c 100644 --- a/Solverz/equation/eqn.py +++ b/Solverz/equation/eqn.py @@ -13,6 +13,7 @@ from Solverz.sym_algebra.matrix_calculus import MixedEquationDiff from Solverz.sym_algebra.transform import finite_difference, semi_descritize from Solverz.num_api.custom_function import numerical_interface +from Solverz.variable.ssymbol import sSym2Sym def sVar2Var(var: Union[Var, iVar, List[iVar, Var]]) -> Union[iVar, List[iVar]]: @@ -55,7 +56,9 @@ def obtain_symbols(self) -> Dict[str, Symbol]: temp_dict[symbol_.name0] = symbol_.symbol0 temp_dict.update(symbol_.SymInIndex) - return temp_dict + # to sort in lexicographic order + sorted_dict = {key: temp_dict[key] for key in sorted(temp_dict)} + return sorted_dict def lambdify(self) -> Callable: return splambdify(self.SYMBOLS.values(), self.RHS, [numerical_interface, 'numpy']) @@ -148,8 +151,8 @@ class Ode(Eqn): def __init__(self, name: str, f, diff_var: Union[iVar, IdxVar, Var]): - super().__init__(name, sVar2Var(f)) - diff_var = sVar2Var(diff_var) + super().__init__(name, sSym2Sym(f)) + diff_var = sSym2Sym(diff_var) self.diff_var = diff_var self.LHS = Derivative(diff_var, t) diff --git a/Solverz/equation/equations.py b/Solverz/equation/equations.py index f55d3cd..ffa7e7d 100644 --- a/Solverz/equation/equations.py +++ b/Solverz/equation/equations.py @@ -17,6 +17,7 @@ from Solverz.variable.variables import Vars from Solverz.utilities.address import Address, combine_Address from Solverz.num_api.Array import Array +from Solverz.equation.jac import Jac, JacBlock class Equations: @@ -39,6 +40,8 @@ def __init__(self, self.PARAM: Dict[str, ParamBase] = dict() self.triggerable_quantity: Dict[str, str] = dict() self.jac_element_address = Address() + self.jac: Jac = Jac() + self.nstep = 0 if isinstance(eqn, Eqn): eqn = [eqn] @@ -58,17 +61,55 @@ def add_eqn(self, eqn: Eqn): self.f_list = self.f_list + [eqn.name] for symbol_ in eqn.SYMBOLS.values(): - if isinstance(symbol_, (Para, iAliasVar)): + if isinstance(symbol_, Para): # this is not fully initialize of Parameters, please use param_initializer - self.PARAM[symbol_.name] = Param(symbol_.name, value=symbol_.value, dim=symbol_.dim) + self.PARAM[symbol_.name] = Param(symbol_.name, + value=symbol_.value, + dim=symbol_.dim) + elif isinstance(symbol_, iAliasVar): + self.PARAM[symbol_.name] = Param(symbol_.name, + value=symbol_.value, + dim=symbol_.dim, + is_alias=True) elif isinstance(symbol_, idx): self.PARAM[symbol_.name] = IdxParam(symbol_.name, value=symbol_.value) self.EQNs[eqn.name].derive_derivative() - def assign_eqn_var_address(self, *xys: Vars): + def assign_eqn_var_address(self, *args): pass + def Fy(self, y) -> List[Tuple[str, str, EqnDiff, np.ndarray]]: + pass + + def FormJac(self, y): + self.assign_eqn_var_address(y) + + Fy_list = self.Fy(y) + for fy in Fy_list: + EqnName = fy[0] + EqnAddr = self.a[EqnName] + VarName = fy[1] + VarAddr = self.var_address[VarName] + DiffVar = fy[2].diff_var + DeriExpr = fy[2].RHS + + DiffVarEqn = Eqn('DiffVarEqn' + DiffVar.name, DiffVar) + args = self.obtain_eqn_args(DiffVarEqn, y, 0) + DiffVarValue = Array(DiffVarEqn.NUM_EQN(*args), dim=1) + + # The value of deri can be either matrix, vector, or scalar(number). We cannot reshape it. + Value0 = np.array(fy[3]) + + jb = JacBlock(EqnName, + EqnAddr, + DiffVar, + DiffVarValue, + VarAddr, + DeriExpr, + Value0) + self.jac.add_block(EqnName, DiffVar, jb) + @property def eqn_size(self): # total size of all the equations @@ -121,38 +162,36 @@ def eval(self, eqn_name: str, *args: Union[np.ndarray]) -> np.ndarray: """ return Array(self.EQNs[eqn_name].NUM_EQN(*args), dim=1) - def trigger_param_updater(self, eqn: Eqn, *xys): + def trigger_param_updater(self, y): # update/initialize triggerable params - if len(xys) > 0: - for para_name, trigger_var in self.triggerable_quantity.items(): - trigger_func = self.PARAM[para_name].trigger_fun - args = [] - for var in trigger_var: - var_value = None - if var in self.PARAM: - var_value = self.PARAM[var].v - else: - for y in xys: - if var in y.var_list: - var_value = y[var] - if var_value is None: - raise ValueError(f'Para/iVar {var} not defined') - else: - args.append(var_value) - temp = trigger_func(*args) - if self.PARAM[para_name].v is None: - self.PARAM[para_name].v = temp + for para_name, trigger_var in self.triggerable_quantity.items(): + trigger_func = self.PARAM[para_name].trigger_fun + args = [] + for var in trigger_var: + var_value = None + if var in self.PARAM: + var_value = self.PARAM[var].v else: - if type(temp) is not type(self.PARAM[para_name].v): - raise TypeError( - f"The return types of trigger func for param {para_name} must be {type(self.PARAM[para_name].v)}") + if var in y.var_list: + var_value = y[var] + if var_value is None: + raise ValueError(f'Para/iVar {var} not defined') + else: + args.append(var_value) + temp = trigger_func(*args) + if self.PARAM[para_name].v is None: + self.PARAM[para_name].v = temp + else: + if type(temp) is not type(self.PARAM[para_name].v): + raise TypeError( + f"The return types of trigger func for param {para_name} must be {type(self.PARAM[para_name].v)}") - def obtain_eqn_args(self, eqn: Eqn, t=None, *xys: Vars) -> List[np.ndarray]: + def obtain_eqn_args(self, eqn: Eqn, y: Vars, t=0) -> List[np.ndarray]: """ Obtain the args of equations """ - self.trigger_param_updater(eqn, *xys) + self.trigger_param_updater(y) args = [] for symbol in eqn.SYMBOLS.values(): @@ -164,10 +203,9 @@ def obtain_eqn_args(self, eqn: Eqn, t=None, *xys: Vars) -> List[np.ndarray]: args.append(temp) value_obtained = True else: - for y in xys: - if symbol.name in y.var_list: - args.append(y[symbol.name]) - value_obtained = True + if symbol.name in y.var_list: + args.append(y[symbol.name]) + value_obtained = True if not value_obtained: raise ValueError(f'Cannot find the values of variable {symbol.name}') return args @@ -182,34 +220,8 @@ def eval_diffs(self, eqn_name: str, var_name: str, *args: np.ndarray) -> np.ndar """ return self.EQNs[eqn_name].derivatives[var_name].NUM_EQN(*args) - def parse_address(self, eqn_name, var_name, eqndiff, t=None, *xys): - var_idx = eqndiff.var_idx - var_idx_func = eqndiff.var_idx_func - - equation_address = self.a.v[eqn_name] - if var_idx is None: - variable_address = self.var_address.v[var_name] - elif isinstance(var_idx, (float, int)): - 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): - variable_address = self.var_address.v[var_name][ - var_idx_func.NUM_EQN(*self.obtain_eqn_args(var_idx_func, t, *xys))] - elif isinstance(var_idx, Expr): - args = self.obtain_eqn_args(var_idx_func, *xys) - variable_address = self.var_address.v[var_name][var_idx_func.NUM_EQN(*args).reshape(-1, )] - elif isinstance(var_idx, list): - variable_address = self.var_address.v[var_name][var_idx] - else: - raise TypeError(f"Unsupported variable index {var_idx} for equation {eqn_name}") - - return equation_address, variable_address - - def evalf(self, expr: Expr, t, *xys: Vars) -> np.ndarray: - eqn = Eqn('Solverz evalf temporal equation', expr) - args = self.obtain_eqn_args(eqn, t, *xys) - return eqn.NUM_EQN(*args) + def evalf(self, *args) -> np.ndarray: + pass class AE(Equations): @@ -225,11 +237,10 @@ def __init__(self, if len(self.f_list) > 0: raise ValueError(f'Ode found. This object should be DAE!') - def assign_eqn_var_address(self, *xys: Vars): + def assign_eqn_var_address(self, y: Vars): """ ASSIGN ADDRESSES TO EQUATIONS """ - y = xys[0] temp = 0 for eqn_name in self.EQNs.keys(): @@ -244,35 +255,6 @@ def assign_eqn_var_address(self, *xys: Vars): self.var_address = y.a self.vsize = y.total_size - # assign dim of derivatives, which is indispensable in Jac printer - gy = self.g_y(y) - for gy_tuple in gy: - eqn_name = gy_tuple[0] - var_name = gy_tuple[1] - eqndiff = gy_tuple[2] - value = gy_tuple[3] - if isinstance(value, (np.ndarray, csc_array)): - # We use coo_array here because by default when converting to CSR or CSC format, - # duplicate (i,j) entries will be summed together. - # This facilitates efficient construction of finite element matrices and the like. - if value.ndim == 2: # matrix - self.EQNs[eqn_name].derivatives[eqndiff.diff_var.name].dim = 2 - - elif value.ndim == 1 and value.shape[0] != 1: # vector - self.EQNs[eqn_name].derivatives[eqndiff.diff_var.name].dim = 1 - # self.jac_element_address.add((eqn_name, var_name, eqndiff), value.shape[0]) - elif value.ndim == 1 and value.shape[0] == 1: # scalar in np.ndarray for example array([0.0]) - self.EQNs[eqn_name].derivatives[eqndiff.diff_var.name].dim = 0 - # self.jac_element_address.add((eqn_name, var_name, eqndiff), self.a.size[eqn_name]) - else: - raise ValueError("Unknown derivative value dimension type!") - self.EQNs[eqn_name].derivatives[eqndiff.diff_var.name].v_type = 'array' - elif isinstance(value, (Number, SymNumber)): - self.EQNs[eqn_name].derivatives[eqndiff.diff_var.name].dim = 0 - self.EQNs[eqn_name].derivatives[eqndiff.diff_var.name].v_type = 'Number' - else: - raise ValueError(f"Unknown derivative data type {type(value)}!") - def g(self, y: Vars, eqn: str = None) -> np.ndarray: """ @@ -283,16 +265,16 @@ def g(self, y: Vars, eqn: str = None) -> np.ndarray: temp = [] if not eqn: for eqn_name, eqn_ in self.EQNs.items(): - args = self.obtain_eqn_args(eqn_, None, y) + args = self.obtain_eqn_args(eqn_, y) g_eqny = self.eval(eqn_name, *args) g_eqny = g_eqny.toarray() if isinstance(g_eqny, csc_array) else g_eqny temp.append(g_eqny.reshape(-1, )) return np.hstack(temp) else: - args = self.obtain_eqn_args(self.EQNs[eqn], None, y) + args = self.obtain_eqn_args(self.EQNs[eqn], y) return self.eval(eqn, *args) - def g_y(self, y: Vars, eqn: List[str] = None, var: List[str] = None) -> List[Tuple[str, str, EqnDiff, np.ndarray]]: + def gy(self, y: Vars, eqn: List[str] = None, var: List[str] = None) -> List[Tuple[str, str, EqnDiff, np.ndarray]]: """ generate Jacobian matrices of Eqn object with respect to var object :param y: @@ -312,11 +294,19 @@ def g_y(self, y: Vars, eqn: List[str] = None, var: List[str] = None) -> List[Tup for var_name in var: for key, value in eqn_diffs.items(): if var_name == value.diff_var_name: # f is viewed as f[k] - args = self.obtain_eqn_args(eqn_diffs[key], None, y) + args = self.obtain_eqn_args(eqn_diffs[key], y) temp = self.eval_diffs(eqn_name, key, *args) gy = [*gy, (eqn_name, var_name, eqn_diffs[key], temp)] return gy + def Fy(self, y): + return self.gy(y) + + def evalf(self, expr: Expr, y: Vars) -> np.ndarray: + eqn = Eqn('Solverz evalf temporal equation', expr) + args = self.obtain_eqn_args(eqn, y) + return eqn.NUM_EQN(*args) + def __repr__(self): if not self.eqn_size: return f"Algebraic equation {self.name} with addresses uninitialized" @@ -360,15 +350,20 @@ def __init__(self, if len(self.f_list) == 0: raise ValueError(f'No ODE found. You should initialise AE instead!') - def assign_eqn_var_address(self, *xys: Vars): + def evalf(self, expr: Expr, t, y: Vars) -> np.ndarray: + eqn = Eqn('Solverz evalf temporary equation', expr) + args = self.obtain_eqn_args(eqn, y, t) + return eqn.NUM_EQN(*args) + + def assign_eqn_var_address(self, y: Vars): """ ASSIGN ADDRESSES TO EQUATIONS f and g """ temp = 0 for eqn_name in self.f_list: - feval = self.f(None, *xys, eqn=eqn_name) - lhs_eval = self.eval_lhs(None, *xys, eqn=eqn_name) + feval = self.f(None, y, eqn=eqn_name) + lhs_eval = self.eval_lhs(None, y, eqn=eqn_name) if isinstance(feval, Number): rhs_size = 1 else: @@ -385,7 +380,7 @@ def assign_eqn_var_address(self, *xys: Vars): self.state_num = temp for eqn_name in self.g_list: - geval = self.g(None, *xys, eqn=eqn_name) + geval = self.g(None, y, eqn=eqn_name) if np.max(np.abs(geval)) > 1e-5: warnings.warn( f'Inconsistent initial values for algebraic equation: {eqn_name}, with deviation {np.max(np.abs(geval))}!') @@ -399,100 +394,60 @@ def assign_eqn_var_address(self, *xys: Vars): self.algebra_num = temp - self.state_num - if len(xys) == 1: - self.var_address = xys[0].a - self.vsize = self.var_address.total_size - elif len(xys) == 2: - x = xys[0] - y = xys[1] - if x.total_size != self.state_num: - raise ValueError("Length of input state variable not compatible with state equations!") - if y.total_size != self.algebra_num: - raise ValueError("Length of input algebraic variable not compatible with algebraic equations!") - self.var_address = combine_Address(x.a, y.a) - - for var_name in self.var_address.v.keys(): - if var_name not in self.SYMBOLS: - raise ValueError(f"DAE {self.name} has no variable {var_name}") - - self.vsize: int = self.var_address.total_size - elif len(xys) > 2: - raise ValueError("Accept at most two positional arguments!") - - # assign dim of derivatives, which is indispensable in Jac printer - fg_xy = self.f_xy(0, *xys) - if len(self.g_list) > 0: - fg_xy.extend(self.g_xy(0, *xys)) - for gy_tuple in fg_xy: - eqn_name = gy_tuple[0] - var_name = gy_tuple[1] - eqndiff = gy_tuple[2] - value = gy_tuple[3] - if isinstance(value, (np.ndarray, csc_array)): - # We use coo_array here because by default when converting to CSR or CSC format, - # duplicate (i,j) entries will be summed together. - # This facilitates efficient construction of finite element matrices and the like. - if value.ndim == 2: # matrix - self.EQNs[eqn_name].derivatives[eqndiff.diff_var.name].dim = 2 - - elif value.ndim == 1 and value.shape[0] != 1: # vector - self.EQNs[eqn_name].derivatives[eqndiff.diff_var.name].dim = 1 - # self.jac_element_address.add((eqn_name, var_name, eqndiff), value.shape[0]) - elif value.ndim == 1 and value.shape[0] == 1: # scalar in np.ndarray for example array([0.0]) - self.EQNs[eqn_name].derivatives[eqndiff.diff_var.name].dim = 0 - # self.jac_element_address.add((eqn_name, var_name, eqndiff), self.a.size[eqn_name]) - else: - raise ValueError("Unknown derivative value dimension type!") - self.EQNs[eqn_name].derivatives[eqndiff.diff_var.name].v_type = 'array' - elif isinstance(value, (Number, SymNumber)): - self.EQNs[eqn_name].derivatives[eqndiff.diff_var.name].dim = 0 - self.EQNs[eqn_name].derivatives[eqndiff.diff_var.name].v_type = 'Number' - else: - raise ValueError(f"Unknown derivative data type {type(value)}!") + self.var_address = y.a + self.vsize = self.var_address.total_size - def F(self, t, *xys) -> np.ndarray: + def F(self, t, y) -> np.ndarray: """ Return [f(t,x,y), g(t,x,y)] :param t: time - :param xys: Vars + :param y: Vars :return: """ if len(self.g_list) > 0: - return np.concatenate([self.f(t, *xys), self.g(t, *xys)]) + return np.concatenate([self.f(t, y), self.g(t, y)]) else: - return self.f(t, *xys) + return self.f(t, y) - def f(self, t, *xys, eqn=None) -> np.ndarray: + def f(self, t, y, eqn=None) -> np.ndarray: temp = [] if eqn: if eqn in self.f_list: - args = self.obtain_eqn_args(self.EQNs[eqn], t, *xys) + args = self.obtain_eqn_args(self.EQNs[eqn], y, t) temp.append(self.eval(eqn, *args).reshape(-1, )) else: for eqn in self.f_list: - args = self.obtain_eqn_args(self.EQNs[eqn], t, *xys) + args = self.obtain_eqn_args(self.EQNs[eqn], y, t) temp.append(self.eval(eqn, *args).reshape(-1, )) return np.hstack(temp) - def eval_lhs(self, t, *xys, eqn=None) -> np.ndarray: + def eval_lhs(self, t, y, eqn=None) -> np.ndarray: temp = [] if eqn: if eqn in self.f_list: - lhs_eqn = Eqn('lhs_' + eqn, self.EQNs[eqn].diff_var) - args = self.obtain_eqn_args(lhs_eqn, t, *xys) - temp.append(Array(lhs_eqn.NUM_EQN(*args), dim=1)) + ode = self.EQNs[eqn] + if isinstance(ode, Ode): + lhs_eqn = Eqn('lhs_' + eqn, ode.diff_var) + args = self.obtain_eqn_args(lhs_eqn, y, t) + temp.append(Array(lhs_eqn.NUM_EQN(*args), dim=1)) + else: + raise TypeError(f"Equation {ode.name} in f_list is not Ode!") else: for eqn in self.f_list: - lhs_eqn = Eqn('lhs_' + eqn, self.EQNs[eqn].diff_var) - args = self.obtain_eqn_args(lhs_eqn, t, *xys) - temp.append(Array(lhs_eqn.NUM_EQN(*args), dim=1)) + ode = self.EQNs[eqn] + if isinstance(ode, Ode): + lhs_eqn = Eqn('lhs_' + eqn, ode.diff_var) + args = self.obtain_eqn_args(lhs_eqn, y, t) + temp.append(Array(lhs_eqn.NUM_EQN(*args), dim=1)) + else: + raise TypeError(f"Equation {ode.name} in f_list is not Ode!") return np.hstack(temp) - def g(self, t, *xys, eqn=None) -> np.ndarray: + def g(self, t, y, eqn=None) -> np.ndarray: """ `xys` is either: @@ -507,61 +462,61 @@ def g(self, t, *xys, eqn=None) -> np.ndarray: temp = [] if eqn: if eqn in self.g_list: - args = self.obtain_eqn_args(self.EQNs[eqn], t, *xys) + args = self.obtain_eqn_args(self.EQNs[eqn], y, t) temp.append(self.eval(eqn, *args).reshape(-1, )) else: for eqn in self.g_list: - args = self.obtain_eqn_args(self.EQNs[eqn], t, *xys) + args = self.obtain_eqn_args(self.EQNs[eqn], y, t) temp.append(self.eval(eqn, *args).reshape(-1, )) return np.hstack(temp) - def f_xy(self, t, *xys: Vars) -> List[Tuple[str, str, EqnDiff, np.ndarray]]: + def fy(self, t, y: Vars) -> List[Tuple[str, str, EqnDiff, np.ndarray]]: """ - generate partial derivatives of f w.r.t. vars in xys - :return: List[Tuple[Equation_name, var_name, np.ndarray]] + generate partial derivatives of f w.r.t. y """ - fxy: List[Tuple[str, str, EqnDiff, np.ndarray]] = [] + fy: List[Tuple[str, str, EqnDiff, np.ndarray]] = [] - var: List[str] = list() - for xy in xys: - var = var + xy.var_list + var: List[str] = y.var_list for eqn_name in self.f_list: eqn_diffs: Dict[str, EqnDiff] = self.EQNs[eqn_name].derivatives for var_name in var: for key, value in eqn_diffs.items(): if var_name == value.diff_var_name: - args = self.obtain_eqn_args(eqn_diffs[key], t, *xys) + args = self.obtain_eqn_args(eqn_diffs[key], y, t) temp = self.eval_diffs(eqn_name, key, *args) - fxy = [*fxy, (eqn_name, var_name, eqn_diffs[key], temp)] - return fxy + fy = [*fy, (eqn_name, var_name, eqn_diffs[key], temp)] + return fy - def g_xy(self, t, *xys: Vars) -> List[Tuple[str, str, EqnDiff, np.ndarray]]: + def gy(self, t, y: Vars) -> List[Tuple[str, str, EqnDiff, np.ndarray]]: """ - generate partial derivatives of f w.r.t. vars in xys - :return: List[Tuple[Equation_name, var_name, np.ndarray]] + generate partial derivatives of g w.r.t. y """ if len(self.g_list) == 0: raise ValueError(f'No AE found in {self.name}!') - gxy: List[Tuple[str, str, EqnDiff, np.ndarray]] = [] + gy: List[Tuple[str, str, EqnDiff, np.ndarray]] = [] - var: List[str] = list() - for xy in xys: - var = var + xy.var_list + var: List[str] = y.var_list for eqn_name in self.g_list: eqn_diffs: Dict[str, EqnDiff] = self.EQNs[eqn_name].derivatives for var_name in var: for key, value in eqn_diffs.items(): if var_name == value.diff_var_name: - args = self.obtain_eqn_args(eqn_diffs[key], t, *xys) + args = self.obtain_eqn_args(eqn_diffs[key], y, t) temp = self.eval_diffs(eqn_name, key, *args) - gxy = [*gxy, (eqn_name, var_name, eqn_diffs[key], temp)] - return gxy + gy = [*gy, (eqn_name, var_name, eqn_diffs[key], temp)] + return gy + + def Fy(self, y): + fg_xy = self.fy(0, y) + if len(self.g_list) > 0: + fg_xy.extend(self.gy(0, y)) + return fg_xy @property def M(self): @@ -613,7 +568,8 @@ def M(self): eqn_address_list = equation_address.tolist() var_address_list = variable_address.tolist() if len(eqn_address_list) != len(var_address_list): - raise ValueError(f"Incompatible eqn address length {len(eqn_address_list)} and variable address length {len(var_address_list)}") + raise ValueError( + f"Incompatible eqn address length {len(eqn_address_list)} and variable address length {len(var_address_list)}") row.extend(eqn_address_list) col.extend(var_address_list) else: diff --git a/Solverz/equation/jac.py b/Solverz/equation/jac.py new file mode 100644 index 0000000..895979b --- /dev/null +++ b/Solverz/equation/jac.py @@ -0,0 +1,293 @@ +from __future__ import annotations +import numpy as np +from typing import Dict, Union, List +import warnings + +from sympy import Expr, Function, Integer +from Solverz.sym_algebra.symbols import iVar, IdxVar +from Solverz.utilities.type_checker import is_vector, is_scalar, is_integer, is_number, PyNumber +from Solverz.sym_algebra.functions import Diag + +SolVar = Union[iVar, IdxVar] + + +class Jac: + + def __init__(self): + self.blocks: Dict[str, Dict[SolVar, JacBlock]] = dict() + + def add_block(self, + eqn_name: str, + diff_var: SolVar, + jb: JacBlock): + if eqn_name in self.blocks: + self.blocks[eqn_name][diff_var] = jb + else: + self.blocks[eqn_name] = dict() + self.blocks[eqn_name][diff_var] = jb + + @property + def JacEleNum(self) -> int: + num = 0 + for jbs_row in self.blocks.values(): + for jb in jbs_row.values(): + if jb.DeriType == 'matrix': + raise NotImplementedError("Matrix derivative type not supported!") + else: + num = num + jb.SpEleSize + return num + + def parse_row_col_data(self): + """ + Parse the row, col and data for sparse coo-jac construction. + """ + 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 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() + col[addr_by_ele] = jb.SpVarAddr.copy() + if jb.IsDeriNumber: + data[addr_by_ele] = jb.DeriExpr + addr_by_ele_0 += jb.SpEleSize + return row, col, data + + +class JacBlock: + + def __init__(self, + EqnName: str, + EqnAddr: slice, + DiffVar, + DiffVarValue, + VarAddr, + DeriExpr: Expr, + Value0: np.ndarray | PyNumber): + """ + var_addr is the address of the non-indexed variable. Fox example, if the diff_var is x[0], then the + var_addr is the address of x. JacBlock will parse the address of x[0] + + jac type: + 0. a column vector + 1. a diagonal matrix + 2. a matrix + """ + self.EqnName = EqnName + self.EqnAddr: slice = EqnAddr + self.DiffVar: SolVar = DiffVar + self.DiffVarValue: np.ndarray = DiffVarValue + self.VarAddr: slice = VarAddr + self.DeriExpr = DeriExpr + self.DeriExprBc = Integer(0) # Derivative broadcast + self.Value0 = Value0 # the value of derivative + self.SpEqnAddr: np.ndarray = np.array([]) + self.SpVarAddr: np.ndarray = np.array([]) + self.SpEleSize = 0 + self.SpDeriExpr: Expr = Integer(0) + self.DenEqnAddr: slice = slice(0) + self.DenVarAddr: slice = slice(0) + self.DenDeriExpr: Expr = Integer(0) + + EqnSize = self.EqnAddr.stop - self.EqnAddr.start + + # find out jac type + if DiffVarValue.size > 1: + DiffVarType = 'vector' + else: + DiffVarType = 'scalar' + + self.DiffVarType = DiffVarType + + if not isinstance(self.Value0, np.ndarray): + if isinstance(self.Value0, PyNumber): + self.Value0 = np.array(self.Value0) + else: + raise TypeError(f"Derivative value type {type(self.Value0)} not supported!") + + if self.Value0.ndim == 2: + DeriType = 'matrix' + elif is_vector(self.Value0): + DeriType = 'vector' + elif is_scalar(self.Value0): + DeriType = 'scalar' + else: + raise TypeError(f"Cant deduce derivative type of value {self.Value0}") + + self.DeriType = DeriType + + # broadcast derivative + if DiffVarType == 'scalar': + match DeriType: + case 'scalar': + self.DeriExprBc = self.DeriExpr * Ones(EqnSize) + case 'vector': + if self.Value0.size != EqnSize: + raise ValueError(f"Vector derivative size {self.Value0.size} != Equation size {EqnSize}") + self.DeriExprBc = self.DeriExpr + case 'matrix': + # self.DeriExprBc = self.DeriExpr + raise TypeError("Matrix derivative of scalar variables not supported!") + case _: + raise TypeError(f"Derivative with value {self.Value0} of scalar variables not supported!") + elif DiffVarType == 'vector': + match DeriType: + case 'scalar': + if self.DiffVarValue.size != EqnSize: + raise ValueError(f"Vector variable {self.DiffVar} size {self.DiffVarValue.size} != " + + f"Equation size {EqnSize} in scalar derivative case.") + self.DeriExprBc = self.DeriExpr * Ones(EqnSize) + case 'vector': + if self.Value0.size != EqnSize: + raise ValueError(f"Vector derivative size {self.Value0.size} != Equation size {EqnSize}") + if self.DiffVarValue.size != EqnSize: + raise ValueError(f"Vector variable {self.DiffVar} size {self.DiffVarValue.size} != " + + f"Equation size {EqnSize} in vector derivative case.") + self.DeriExprBc = self.DeriExpr + case 'matrix': + if self.Value0.shape[0] == EqnSize: + try: + self.Value0 @ DiffVarValue + except ValueError: + raise ValueError(f"Incompatible matrix derivative size {self.Value0.shape} " + + f"and vector variable size {DiffVarValue.shape}.") + self.DeriExprBc = self.DeriExpr + case _: + raise TypeError(f"Derivative with value {self.Value0} of vector variables not supported!") + + # parse sparse jac blocks, the address and expression + self.ParseSp() + # parse dense jac blocks, the address and expression + self.ParseDen() + + @property + def IsDeriNumber(self): + return is_number(self.DeriExpr) + + def ParseSp(self): + EqnSize = self.EqnAddr.stop - self.EqnAddr.start + match self.DiffVarType: + case 'vector': + match self.DeriType: + case 'matrix': + warnings.warn("Sparse parser of matrix type jac block not implemented!") + case 'vector' | 'scalar': + self.SpEqnAddr = slice2array(self.EqnAddr) + if isinstance(self.DiffVar, iVar): + self.SpVarAddr = slice2array(self.VarAddr) + else: + if isinstance(self.DiffVar.index, slice): + VarArange = slice2array(self.VarAddr)[self.DiffVar.index] + self.SpVarAddr = VarArange + elif is_integer(self.DiffVar.index): + raise TypeError(f"Index of vector variable cant be integer!") + else: + raise TypeError(f"Index type {type(self.DiffVar.index)} not supported!") + self.SpDeriExpr = self.DeriExprBc + case 'scalar': + match self.DeriType: + case 'vector' | 'scalar': + self.SpEqnAddr = np.arange(self.EqnAddr.start, self.EqnAddr.stop) + if isinstance(self.DiffVar, iVar): + idx = self.VarAddr.start + self.SpVarAddr = np.array(EqnSize * [idx]).reshape((-1,)).astype(int) + 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: + idx = VarArange[0] + self.SpVarAddr = np.array(EqnSize * [idx]).reshape((-1,)).astype(int) + elif is_integer(self.DiffVar.index): + idx = slice2array(self.VarAddr)[self.DiffVar.index] + self.SpVarAddr = np.array(EqnSize * [idx]).reshape((-1,)).astype(int) + else: + raise TypeError(f"Index type {type(self.DiffVar.index)} not supported!") + self.SpDeriExpr = self.DeriExprBc + if self.SpEqnAddr.size != self.SpVarAddr.size: + raise ValueError(f"Incompatible equation size {self.SpEqnAddr.size} of Equation {self.EqnName} " + + f"and variable size {self.SpVarAddr.size} of Variable {self.DiffVar}!") + self.SpEleSize = self.SpEqnAddr.size + + def ParseDen(self): + self.DenEqnAddr = self.EqnAddr + match self.DiffVarType: + case 'vector': + match self.DeriType: + case 'matrix': + if isinstance(self.DiffVar, iVar): + self.DenVarAddr = self.VarAddr + else: + if isinstance(self.DiffVar.index, slice): + VarArange = slice2array(self.VarAddr)[self.DiffVar.index] + self.DenVarAddr = slice(VarArange[0], VarArange[-1] + 1) + elif is_integer(self.DiffVar.index): + raise TypeError(f"Index of vector variable cant be integer!") + else: + raise TypeError(f"Index type {type(self.DiffVar.index)} not supported!") + self.DenDeriExpr = self.DeriExprBc + case 'vector' | 'scalar': + self.DenEqnAddr = self.EqnAddr + if isinstance(self.DiffVar, iVar): + self.DenVarAddr = self.VarAddr + else: + if isinstance(self.DiffVar.index, slice): + VarArange = slice2array(self.VarAddr)[self.DiffVar.index] + self.DenVarAddr = slice(VarArange[0], VarArange[-1] + 1) + elif is_integer(self.DiffVar.index): + raise TypeError(f"Index of vector variable cant be integer!") + else: + raise TypeError(f"Index type {type(self.DiffVar.index)} not supported!") + self.DenDeriExpr = Diag(self.DeriExprBc) + case 'scalar': + match self.DeriType: + case 'vector' | 'scalar': + self.DenEqnAddr = self.EqnAddr + if isinstance(self.DiffVar, iVar): + self.DenVarAddr = self.VarAddr + 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) + elif is_integer(self.DiffVar.index): + idx = int(slice2array(self.VarAddr)[self.DiffVar.index]) + self.DenVarAddr = slice(idx, idx + 1) + else: + raise TypeError(f"Index type {type(self.DiffVar.index)} not supported!") + self.DenDeriExpr = self.DeriExprBc + + +def slice2array(s: slice) -> np.ndarray: + return np.arange(s.start, s.stop) + + +class Ones(Function): + r""" + The all-one vector to broadcast scalars to a vector + For example, 2 -> 2*Ones(eqn_size) + The derivative is d(Ones(eqn_size))/dx=0 + """ + + @classmethod + def eval(cls, x): + if not isinstance(x, Integer): + raise ValueError("The arg of Ones() should be integer.") + + def _eval_derivative(self, s): + return 0 + + def _numpycode(self, printer, **kwargs): + x = self.args[0] + return r'ones(' + printer._print(x, **kwargs) + r')' + + def _lambdacode(self, printer, **kwargs): + return self._numpycode(printer, **kwargs) + + def _pythoncode(self, printer, **kwargs): + return self._numpycode(printer, **kwargs) diff --git a/Solverz/equation/param.py b/Solverz/equation/param.py index 3caebba..61d4ec4 100644 --- a/Solverz/equation/param.py +++ b/Solverz/equation/param.py @@ -17,7 +17,8 @@ def __init__(self, trigger_fun: Callable = None, dim: int = 1, dtype=float, - sparse=False): + sparse=False, + is_alias=False): self.name = name self.triggerable = triggerable self.trigger_var = [trigger_var] if isinstance(trigger_var, str) else trigger_var @@ -27,6 +28,7 @@ def __init__(self, self.sparse = sparse self.__v = None self.v = value + self.is_alias = is_alias # if the Param is an alias var @property def v(self): @@ -57,7 +59,8 @@ def __init__(self, trigger_fun: Callable = None, dim: int = 1, dtype=float, - sparse=False + sparse=False, + is_alias=False ): ParamBase.__init__(self, name, @@ -67,7 +70,8 @@ def __init__(self, trigger_fun, dim, dtype, - sparse) + sparse, + is_alias) sSymBasic.__init__(self, name=name, Type='Para', value=value, dim=dim) diff --git a/Solverz/equation/test/test_jac.py b/Solverz/equation/test/test_jac.py index 118707b..454302f 100644 --- a/Solverz/equation/test/test_jac.py +++ b/Solverz/equation/test/test_jac.py @@ -1,65 +1,376 @@ 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.variable.variables import combine_Vars, as_Vars +from Solverz.equation.jac import JacBlock, Ones, Jac +from Solverz.sym_algebra.functions import Diag def test_jac(): - x = iVar('x', 1) - y = iVar('y', 1) - - f = Ode(name='f', f=-x ** 3 + 0.5 * y ** 2, diff_var=x) - g = Eqn(name='g', eqn=x ** 2 + y ** 2 - 2) - dae = DAE([f, g]) - - z = combine_Vars(as_Vars(x), as_Vars(y)) - fxy = dae.f_xy(None, z) - gxy = dae.g_xy(None, z) - assert fxy[0][3].shape == (1,) - assert np.all(np.isclose(fxy[0][3], [-3.])) - assert fxy[1][3].shape == (1,) - assert np.all(np.isclose(fxy[1][3], [1])) - assert gxy[0][3].shape == (1,) - assert np.all(np.isclose(gxy[0][3], [2])) - assert gxy[1][3].shape == (1,) - assert np.all(np.isclose(gxy[1][3], [2])) - - x = iVar('x', [1, 2, 3]) - i = idx('i', value=[0, 2]) - - f = Eqn('f', eqn=x) - ae = AE(f) - gy = ae.g_y(as_Vars(x)) - assert isinstance(gy[0][3], Number) - # assert gy[0][3].ndim == 0 - assert np.all(np.isclose(gy[0][3], 1)) - - f = Eqn('f', eqn=x[i]) - ae = AE(f) - gy = ae.g_y(as_Vars(x)) - assert isinstance(gy[0][3], Number) - # assert gy[0][3].ndim == 0 - assert np.all(np.isclose(gy[0][3], 1)) - - f = Eqn('f', eqn=x[i] ** 2) - ae = AE(f) - gy = ae.g_y(as_Vars(x)) - assert isinstance(gy[0][3], np.ndarray) - assert gy[0][3].ndim == 1 - assert np.all(np.isclose(gy[0][3], [2., 6.])) - - A_v = np.random.rand(3, 3) - A = Para('A', dim=2) - f = Eqn('f', eqn=A * x) - ae = AE(f) - ae.param_initializer('A', Param('A', value=A_v, dim=2)) - gy = ae.g_y(as_Vars(x)) - assert isinstance(gy[0][3], np.ndarray) - assert gy[0][3].ndim == 2 - np.testing.assert_allclose(gy[0][3], A_v) + # test jac element number counter + jac = Jac() + x = iVar('x') + y = iVar('y') + omega = iVar('omega') + jac.add_block('a', + x[0], + JacBlock('a', + slice(0, 3), + x[0], + np.array([1]), + slice(1, 2), + iVar('y'), + np.array([1]))) + jac.add_block('b', + omega[1:4], + JacBlock('b', + slice(3, 12), + omega[4:13], + np.ones(9), + slice(3, 100), + iVar('y') ** 2, + np.ones(9))) + jac.add_block('b', + omega[5:7], + JacBlock('b', + slice(0, 3), + omega[13], + np.ones(1), + slice(3, 100), + Integer(5), + 5)) + assert jac.JacEleNum == 15 + row, col, data = jac.parse_row_col_data() + assert row.size == 15 + assert col.size == 15 + assert data.size == 15 + np.testing.assert_allclose(row, np.concatenate([np.arange(0, 12), np.array([0, 1, 2])])) + np.testing.assert_allclose(col, np.concatenate([np.array([1, 1, 1]), np.arange(7, 16), np.array([16, 16, 16])])) + np.testing.assert_allclose(data, np.concatenate([np.zeros(12), np.array([5, 5, 5])])) + + +#%% scalar var and scalar derivative +def test_jb_scalar_var_scalar_deri(): + # non-index var + jb = JacBlock('a', + slice(0, 3), + iVar('x'), + np.array([1]), + slice(1, 2), + iVar('y'), + np.array([1])) + + assert jb.DenEqnAddr == slice(0, 3) + assert jb.DenVarAddr == slice(1, 2) + 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])) + assert jb.SpDeriExpr == iVar('y') * Ones(3) + assert jb.SpEleSize == 3 + assert jb.IsDeriNumber is False + + # indexed var + jb = JacBlock('a', + slice(0, 3), + iVar('x')[1], + np.array([1]), + slice(1, 3), + iVar('y'), + np.array([1])) + + assert jb.DenEqnAddr == slice(0, 3) + assert jb.DenVarAddr == slice(2, 3) + 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])) + assert jb.SpDeriExpr == iVar('y') * Ones(3) + assert jb.SpEleSize == 3 + assert jb.IsDeriNumber is False + + # sliced var + jb = JacBlock('a', + slice(0, 3), + iVar('x')[1:2], + np.array([1]), + slice(1, 3), + iVar('y'), + np.array([1])) + + assert jb.DenEqnAddr == slice(0, 3) + assert jb.DenVarAddr == slice(2, 3) + 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])) + assert jb.SpDeriExpr == iVar('y') * Ones(3) + assert jb.SpEleSize == 3 + assert jb.IsDeriNumber is False + + +# %% scalar var and vector derivative +def test_jb_scalar_var_vector_deri(): + # non-index var + jb = JacBlock('a', + slice(0, 3), + iVar('x'), + np.array([1]), + slice(1, 2), + iVar('y'), + np.array([1, 1, 1])) + + assert jb.DenEqnAddr == slice(0, 3) + assert jb.DenVarAddr == slice(1, 2) + assert jb.DenDeriExpr == iVar('y') + assert_allclose(jb.SpEqnAddr, np.array([0, 1, 2])) + assert_allclose(jb.SpVarAddr, np.array([1, 1, 1])) + assert jb.SpDeriExpr == iVar('y') + assert jb.SpEleSize == 3 + assert jb.IsDeriNumber is False + + # indexed var + jb = JacBlock('a', + slice(0, 3), + iVar('x')[1], + np.array([1]), + slice(1, 3), + iVar('y'), + np.array([1, 1, 1])) + + assert jb.DenEqnAddr == slice(0, 3) + assert jb.DenVarAddr == slice(2, 3) + assert jb.DenDeriExpr == iVar('y') + assert_allclose(jb.SpEqnAddr, np.array([0, 1, 2])) + assert_allclose(jb.SpVarAddr, np.array([2, 2, 2])) + assert jb.SpDeriExpr == iVar('y') + assert jb.SpEleSize == 3 + assert jb.IsDeriNumber is False + + # sliced var + jb = JacBlock('a', + slice(0, 3), + iVar('x')[1:2], + np.array([1]), + slice(1, 3), + iVar('y'), + np.array([1, 1, 1])) + + assert jb.DenEqnAddr == slice(0, 3) + assert jb.DenVarAddr == slice(2, 3) + assert jb.DenDeriExpr == iVar('y') + assert_allclose(jb.SpEqnAddr, np.array([0, 1, 2])) + assert_allclose(jb.SpVarAddr, np.array([2, 2, 2])) + assert jb.SpDeriExpr == iVar('y') + assert jb.SpEleSize == 3 + assert jb.IsDeriNumber is False + + # Derivative size not compatible with equation size + with pytest.raises(ValueError, match="Vector derivative size 2 != Equation size 3"): + jb = JacBlock('a', + slice(0, 3), + iVar('x')[1:2], + np.array([1]), + slice(1, 3), + iVar('y'), + np.array([1, 1])) + + +# %% vector var and scalar derivative +def test_jb_vector_var_scalar_deri(): + # non-index var + jb = JacBlock('a', + slice(1, 10), + iVar('x'), + np.ones(9), + slice(1, 10), + iVar('y'), + np.array([1])) + + assert jb.DenEqnAddr == slice(1, 10) + assert jb.DenVarAddr == slice(1, 10) + assert jb.DenDeriExpr == Diag(iVar('y') * Ones(9)) + assert_allclose(jb.SpEqnAddr, np.arange(1, 10)) + assert_allclose(jb.SpVarAddr, np.arange(1, 10)) + assert jb.SpDeriExpr == iVar('y') * Ones(9) + assert jb.SpEleSize == 9 + assert jb.IsDeriNumber is False + + # number as derivative + jb = JacBlock('a', + slice(1, 10), + iVar('x'), + np.ones(9), + slice(1, 10), + Integer(2), + 2) + + assert jb.DenEqnAddr == slice(1, 10) + assert jb.DenVarAddr == slice(1, 10) + assert jb.DenDeriExpr == Diag(2 * Ones(9)) + assert_allclose(jb.SpEqnAddr, np.arange(1, 10)) + assert_allclose(jb.SpVarAddr, np.arange(1, 10)) + assert jb.SpDeriExpr == 2 * Ones(9) + assert jb.SpEleSize == 9 + assert jb.IsDeriNumber is True + + jb = JacBlock('a', + slice(1, 10), + iVar('x'), + np.ones(9), + slice(1, 10), + Integer(2), + np.array([2])) + + assert jb.DenEqnAddr == slice(1, 10) + assert jb.DenVarAddr == slice(1, 10) + assert jb.DenDeriExpr == Diag(2 * Ones(9)) + assert_allclose(jb.SpEqnAddr, np.arange(1, 10)) + assert_allclose(jb.SpVarAddr, np.arange(1, 10)) + assert jb.SpDeriExpr == 2 * Ones(9) + assert jb.SpEleSize == 9 + assert jb.IsDeriNumber is True + + jb = JacBlock('a', + slice(1, 10), + iVar('x'), + np.ones(9), + slice(1, 10), + Integer(2), + np.array(2)) + + assert jb.DenEqnAddr == slice(1, 10) + assert jb.DenVarAddr == slice(1, 10) + assert jb.DenDeriExpr == Diag(2 * Ones(9)) + assert_allclose(jb.SpEqnAddr, np.arange(1, 10)) + assert_allclose(jb.SpVarAddr, np.arange(1, 10)) + assert jb.SpDeriExpr == 2 * Ones(9) + assert jb.SpEleSize == 9 + assert jb.IsDeriNumber is True + + # indexed var + with pytest.raises(TypeError, + match=re.escape("Index of vector variable cant be integer!")): + jb = JacBlock('a', + slice(0, 3), + iVar('x')[1], + np.array([2, 3, 4]), + slice(10, 20), + iVar('y')[0], + np.array([1])) + + # sliced var + jb = JacBlock('a', + slice(0, 3), + iVar('x')[1:4], + np.array([2, 3, 4]), + slice(10, 20), + iVar('y')[0], + np.array([1])) + assert jb.DenEqnAddr == slice(0, 3) + assert jb.DenVarAddr == slice(11, 14) + assert jb.DenDeriExpr == Diag(iVar('y')[0] * Ones(3)) + assert_allclose(jb.SpEqnAddr, np.array([0, 1, 2])) + assert_allclose(jb.SpVarAddr, np.array([11, 12, 13])) + assert jb.SpDeriExpr == iVar('y')[0] * Ones(3) + assert jb.SpEleSize == 3 + assert jb.IsDeriNumber is False + + # sliced var with incompatible eqn and var size + with pytest.raises(ValueError, + match=re.escape("Vector variable x[1:3] size 2 != Equation size 3 in scalar derivative case.")): + jb = JacBlock('a', + slice(0, 3), + iVar('x')[1:3], + np.array([2, 3]), + slice(1, 4), + iVar('y'), + np.array([1])) + + +# %% vector var and vector derivative +def test_jb_vector_var_vector_deri(): + jb = JacBlock('a', + slice(1, 10), + iVar('x'), + np.ones(9), + slice(1, 10), + iVar('y'), + np.ones(9)) + + assert jb.DenEqnAddr == slice(1, 10) + assert jb.DenVarAddr == slice(1, 10) + assert jb.DenDeriExpr == Diag(iVar('y')) + assert_allclose(jb.SpEqnAddr, np.arange(1, 10)) + assert_allclose(jb.SpVarAddr, np.arange(1, 10)) + assert jb.SpDeriExpr == iVar('y') + assert jb.SpEleSize == 9 + assert jb.IsDeriNumber is False + + # incompatible eqn and var size + with pytest.raises(ValueError, + match=re.escape("Vector variable x[1:3] size 2 != Equation size 3 in vector derivative case.")): + jb = JacBlock('a', + slice(0, 3), + iVar('x')[1:3], + np.array([2, 3]), + slice(1, 4), + iVar('y'), + np.array([1, 1, 1])) + + # incompatible eqn and derivative size + with pytest.raises(ValueError, + match=re.escape("Vector derivative size 4 != Equation size 3")): + jb = JacBlock('a', + slice(0, 3), + iVar('x')[1:4], + np.array([2, 3, 4]), + slice(1, 4), + iVar('y'), + np.array([1, 1, 1, 1])) + + +# %% scalar var and matrix derivative +def test_jb_vector_var_matrix_deri(): + with pytest.raises(TypeError, + match=re.escape("Matrix derivative of scalar variables not supported!")): + jb = JacBlock('a', + slice(0, 3), + iVar('x')[1], + np.array([2]), + slice(1, 4), + iVar('y'), + np.zeros((3, 3))) + + # vector var and matrix derivative + # compatible vector and matrix size + with pytest.warns(UserWarning) as record: + jb = JacBlock('a', + slice(0, 3), + iVar('x'), + np.array([2, 3, 4]), + slice(1, 4), + iVar('y'), + np.zeros((3, 3))) + + assert str(record[0].message) == "Sparse parser of matrix type jac block not implemented!" + assert jb.DenEqnAddr == slice(0, 3) + assert jb.DenVarAddr == slice(1, 4) + # incompatible vector and matrix size + with pytest.raises(ValueError, + match=re.escape("Incompatible matrix derivative size (3, 4) and vector variable size (3,).")): + jb = JacBlock('a', + slice(0, 3), + iVar('x'), + np.array([2, 3, 4]), + slice(1, 4), + iVar('y'), + np.zeros((3, 4))) diff --git a/Solverz/model/basic.py b/Solverz/model/basic.py index e4fa66c..ea0ff2c 100644 --- a/Solverz/model/basic.py +++ b/Solverz/model/basic.py @@ -85,7 +85,7 @@ def create_instance(self): if eqn_type == 'FDAE': for i in range(nstep): eqs.update_param(y0.derive_alias(f'_tag_{i}')) - eqs.assign_eqn_var_address(y0) + eqs.FormJac(y0) if eqs.eqn_size != eqs.vsize: warnings.warn(f'Equation size {eqs.eqn_size} and variable size {eqs.vsize} not equal!') diff --git a/Solverz/sym_algebra/functions.py b/Solverz/sym_algebra/functions.py index 308f125..23492c3 100644 --- a/Solverz/sym_algebra/functions.py +++ b/Solverz/sym_algebra/functions.py @@ -182,6 +182,7 @@ class Abs(UniVarFunc): \frac{\mathrm{d}}{\mathrm{d}x}{\operatorname{Abs}}(x)=\operatorname{Sign}(x). """ + def fdiff(self, argindex=1): """ Get the first derivative of the argument to Abs(). @@ -202,6 +203,7 @@ class exp(UniVarFunc): r""" The exponential function, $e^x$. """ + def fdiff(self, argindex=1): return exp(*self.args) @@ -216,6 +218,7 @@ class sin(UniVarFunc): r""" The sine function. """ + def fdiff(self, argindex=1): if argindex == 1: return Symcos(self.args[0]) @@ -233,6 +236,7 @@ class cos(UniVarFunc): r""" The cosine function. """ + def fdiff(self, argindex=1): if argindex == 1: return -Symsin(self.args[0]) @@ -260,6 +264,7 @@ class Sign(UniVarFunc): \end{cases} """ + def fdiff(self, argindex=1): # sign function should be treated as a constant. if argindex == 1: @@ -320,6 +325,7 @@ class Saturation(MulVarFunc): v_\min&v< v_\min \end{cases} """ + @classmethod def eval(cls, v, vmin, vmax): return v * In(v, vmin, vmax) + vmax * GreaterThan(v, vmax) + vmin * LessThan(v, vmin) @@ -378,6 +384,7 @@ class Min(MulVarFunc): \end{cases}\end{split} """ + @classmethod def eval(cls, x, y): return x * LessThan(x, y) + y * (1 - LessThan(x, y)) @@ -417,7 +424,7 @@ def _eval_derivative(self, s): def _sympystr(self, printer, **kwargs): return '(({op1})>({op2}))'.format(op1=printer._print(self.args[0]), - op2=printer._print(self.args[1])) + op2=printer._print(self.args[1])) def _numpycode(self, printer, **kwargs): return r'SolGreaterThan(' + ', '.join([printer._print(arg, **kwargs) for arg in self.args]) + r')' @@ -439,7 +446,7 @@ def _eval_derivative(self, s): def _sympystr(self, printer, **kwargs): return '(({op1})<({op2}))'.format(op1=printer._print(self.args[0]), - op2=printer._print(self.args[1])) + op2=printer._print(self.args[1])) def _numpycode(self, printer, **kwargs): return r'SolLessThan(' + ', '.join([printer._print(arg, **kwargs) for arg in self.args]) + r')' @@ -525,22 +532,6 @@ def eval(cls, *args): raise TypeError(f"minmod takes at most 3 positional arguments but {len(args)} were given!") -class zeros(Function): - # print zeros(6,6) as zeros((6,6)) - # or zeros(6,) as zeros((6,)) - def _numpycode(self, printer, **kwargs): - if len(self.args) == 2: - temp1 = printer._print(self.args[0]) - temp2 = printer._print(self.args[1]) - return r'zeros((' + temp1 + ', ' + temp2 + r'))' - elif len(self.args) == 1: - temp = printer._print(self.args[0]) - return r'zeros((' + temp + ', ' + r'))' - - def _pythoncode(self, printer, **kwargs): - return self._numpycode(printer, **kwargs) - - class CSC_array(Function): @classmethod diff --git a/Solverz/utilities/type_checker.py b/Solverz/utilities/type_checker.py index 6a60105..d459962 100644 --- a/Solverz/utilities/type_checker.py +++ b/Solverz/utilities/type_checker.py @@ -1,12 +1,34 @@ -from sympy import Number as symNumber +from sympy import Number as SymNumber, Integer import numpy as np -from numbers import Number +from numbers import Number as PyNumber def is_number(num): - if isinstance(num, (Number, symNumber)): + if isinstance(num, (PyNumber, SymNumber)): return True elif isinstance(num, np.ndarray) and num.size == 1: return True else: return False + + +def is_integer(num): + if isinstance(num, (int, Integer)): + return True + else: + return False + + +def is_vector(a: np.ndarray): + if a.ndim == 1 and a.size > 1: + return True + + +def is_scalar(a): + if is_number(a): + return True + elif isinstance(a, np.ndarray): + if a.ndim == 1 and a.size == 1: + return True + else: + return False diff --git a/instances/4node3pipe.xlsx b/instances/4node3pipe.xlsx deleted file mode 100644 index 51fc417..0000000 Binary files a/instances/4node3pipe.xlsx and /dev/null differ diff --git a/instances/4node3pipe_bench.xlsx b/instances/4node3pipe_bench.xlsx deleted file mode 100644 index 53d9814..0000000 Binary files a/instances/4node3pipe_bench.xlsx and /dev/null differ diff --git a/instances/4node3pipe_change_sign.xlsx b/instances/4node3pipe_change_sign.xlsx deleted file mode 100644 index b503aba..0000000 Binary files a/instances/4node3pipe_change_sign.xlsx and /dev/null differ diff --git a/instances/4node3pipe_change_sign_bench.xlsx b/instances/4node3pipe_change_sign_bench.xlsx deleted file mode 100644 index 957b9f3..0000000 Binary files a/instances/4node3pipe_change_sign_bench.xlsx and /dev/null differ diff --git a/instances/dynamic_gas_flow_single_pipe.xlsx b/instances/dynamic_gas_flow_single_pipe.xlsx deleted file mode 100644 index 9ef1b20..0000000 Binary files a/instances/dynamic_gas_flow_single_pipe.xlsx and /dev/null differ diff --git a/instances/ode_test.xlsx b/instances/ode_test.xlsx deleted file mode 100644 index 0135686..0000000 Binary files a/instances/ode_test.xlsx and /dev/null differ diff --git a/instances/test_burger.xlsx b/instances/test_burger.xlsx deleted file mode 100644 index ec808cd..0000000 Binary files a/instances/test_burger.xlsx and /dev/null differ diff --git a/instances/test_m3b9.xlsx b/instances/test_m3b9.xlsx deleted file mode 100644 index f5191c6..0000000 Binary files a/instances/test_m3b9.xlsx and /dev/null differ diff --git a/instances/test_ode45.xlsx b/instances/test_ode45.xlsx deleted file mode 100644 index bb07457..0000000 Binary files a/instances/test_ode45.xlsx and /dev/null differ diff --git a/instances/test_sirk.xlsx b/instances/test_sirk.xlsx deleted file mode 100644 index 6ba363d..0000000 Binary files a/instances/test_sirk.xlsx and /dev/null differ diff --git a/pyproject.toml b/pyproject.toml index 82f4b8b..c844c18 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,10 +1,13 @@ [build-system] -requires = ["hatchling"] -build-backend = "hatchling.build" +requires = ["setuptools>=64", "setuptools_scm>=8"] +build-backend = "setuptools.build_meta" + +[tool.setuptools_scm] +version_file = "Solverz/_version.py" [project] name = "Solverz" -version = "0.0.1rc4" +dynamic = ["version"] dependencies = [ "sympy>=1.11.1", "numba == 0.58.1", @@ -16,6 +19,7 @@ dependencies = [ "matplotlib>=3.5.2", "tqdm>=4.64.1", "pytest>=7.2.2", + "pytest-datadir>=1.5.0", "networkx >= 3.1" ] authors = [ diff --git a/instances/dae_test.xlsx b/tests/dae_test.xlsx similarity index 100% rename from instances/dae_test.xlsx rename to tests/dae_test.xlsx diff --git a/tests/test_dae.py b/tests/test_dae.py index 676a3aa..dc9af13 100644 --- a/tests/test_dae.py +++ b/tests/test_dae.py @@ -13,7 +13,7 @@ z = as_Vars([x, y]) ndae = made_numerical(dae, z) -df = pd.read_excel('instances/dae_test.xlsx', +df = pd.read_excel('tests/dae_test.xlsx', sheet_name=None, engine='openpyxl' ) diff --git a/tests/test_dhspf.py b/tests/test_dhspf.py deleted file mode 100644 index a0e9774..0000000 --- a/tests/test_dhspf.py +++ /dev/null @@ -1,112 +0,0 @@ -# import pandas as pd -# import numpy as np -# -# from Solverz import Eqn, AE, nr_method, as_Vars, iVar, Para, idx, Abs, exp, Mat_Mul, made_numerical, Param, IdxParam -# -# # %% initialize variables and params -# sys_df = pd.read_excel('instances/4node3pipe.xlsx', -# sheet_name=None, -# engine='openpyxl', -# index_col=0 -# ) -# -# -# def derive_incidence_matrix(f_node, t_node): -# V = np.zeros((max(max(f_node), max(t_node)) + 1, len(f_node))) -# for pipe in range(len(f_node)): -# V[f_node[pipe], pipe] = -1 -# V[t_node[pipe], pipe] = 1 -# return V -# -# -# def derive_v_plus(v: np.ndarray) -> np.ndarray: -# return (v + np.abs(v)) / 2 -# -# -# def derive_v_minus(v: np.ndarray) -> np.ndarray: -# return (v - np.abs(v)) / 2 -# -# -# # %% md -# # Declare equations and parameters -# # %% -# Node = sys_df['Node'] -# i = idx(name='i', value=np.asarray(Node['type'][Node['type'] == 3].index)) # intermediate nodes -# l = idx(name='l', value=np.asarray(Node['type'][Node['type'] == 2].index)) # load nodes -# s = idx(name='s', value=np.asarray(Node['type'][Node['type'] == 1].index)) # non-balance source nodes -# r = idx(name='r', value=np.asarray(Node['type'][Node['type'] == 0].index)) # balance source nodes -# sl = idx(name='sl', value=np.concatenate((s.value, l.value))) -# rs = idx(name='rs', value=np.concatenate((r.value, s.value))) -# rsl = idx(name='rsl', value=np.concatenate((r.value, s.value, l.value))) -# li = idx(name='li', value=np.concatenate((l.value, i.value))) -# rsi = idx(name='rsi', value=np.concatenate((r.value, s.value, i.value))) -# -# mL = Para('mL', dim=2, value=np.asarray(sys_df['m_L'])) -# K = Para(name='K', value=np.asarray(sys_df['K']).reshape(-1, )) -# Ts_set = Para(name='Ts_set', value=np.asarray(sys_df['Ts_set']).reshape(-1, )) -# Tr_set = Para(name='Tr_set', value=np.asarray(sys_df['Tr_set']).reshape(-1, )) -# phi_set = Para(name='phi_set', value=np.asarray(sys_df['phi_set']).reshape(-1, )) -# -# Cp = Para(name='Cp', value=np.asarray(sys_df['Cp']).reshape(-1, )) -# L = Para(name='L', value=np.asarray(sys_df['L']).reshape(-1, )) -# coeff_lambda = Para(name='coeff_lambda', value=np.asarray(sys_df['coeff_lambda']).reshape(-1, )) -# Ta = Para(name='Ta', value=np.asarray(sys_df['Ta']).reshape(-1, )) -# f_node = sys_df['Pipe']['f_node'] -# t_node = sys_df['Pipe']['t_node'] -# V = Para(name='V', dim=2, value=derive_incidence_matrix(f_node, t_node)) -# Vp = Para(name='Vp', dim=2, value=derive_v_plus(V.value)) -# Vm = Para(name='Vm', dim=2, value=derive_v_minus(V.value)) -# f_node = idx(name='f_node', value=np.asarray(f_node).reshape(-1, )) # intermediate nodes -# t_node = idx(name='t_node', value=np.asarray(t_node).reshape(-1, )) # load nodes -# -# m = iVar(name='m', value=np.asarray(sys_df['var']['m'])) -# mq = iVar(name='mq', value=np.asarray(sys_df['var']['mq'])) -# Ts = iVar(name='Ts', value=np.asarray(sys_df['var']['Ts'])) -# Tr = iVar(name='Tr', value=np.asarray(sys_df['var']['Tr'])) -# Touts = iVar(name='Touts', value=np.asarray(sys_df['var']['Touts'])) -# Toutr = iVar(name='Toutr', value=np.asarray(sys_df['var']['Toutr'])) -# phi = iVar(name='phi', value=np.asarray(sys_df['var']['phi'])) -# -# E1 = Eqn(name='E1', eqn=(Ts[f_node] - Ta) * exp(-coeff_lambda * L / (Cp * Abs(m))) + Ta - Touts) -# E3 = Eqn(name='E3', eqn=(Tr[t_node] - Ta) * exp(-coeff_lambda * L / (Cp * Abs(m))) + Ta - Toutr) -# E5 = Eqn(name='E5', eqn=Mat_Mul(V[rs, :], m) + mq[rs]) -# E6 = Eqn(name='E6', eqn=Mat_Mul(V[l, :], m) - mq[l]) -# E7 = Eqn(name='E7', eqn=Mat_Mul(V[i, :], m)) -# E8 = Eqn(name='E8', eqn=Mat_Mul(mL, K * Abs(m) * m)) -# E9 = Eqn(name='E9', eqn=Ts[li] * Mat_Mul(Vp[li, :], Abs(m)) - Mat_Mul(Vp[li, :], Touts * Abs(m))) -# E10 = Eqn(name='E10', eqn=Tr[rsi] * Mat_Mul(Vm[rsi, :], Abs(m)) - Mat_Mul(Vm[rsi, :], Toutr * Abs(m))) -# E11 = Eqn(name='E11', eqn=phi[rsl] - 4182 * mq[rsl] * (Ts[rsl] - Tr[rsl])) -# E12 = Eqn(name='E12', eqn=Ts[rs] - Ts_set) -# E13 = Eqn(name='E13', eqn=Tr[l] - Tr_set) -# E14 = Eqn(name='E14', eqn=phi[sl] - phi_set) -# E15 = Eqn(name='E15', eqn=phi[i]) -# E16 = Eqn(name='E16', eqn=mq[i]) -# E = AE(name='Pipe Equations', eqn=[E1, E3, E5, E6, E7, E8, E9, E10, E11, E12, E13, E14, E15, E16]) -# y0 = as_Vars([m, mq, Ts, Tr, Touts, Toutr, phi]) -# -# nE, code = made_numerical(E, y0, output_code=True) -# -# sol = nr_method(nE, y0) -# # y_cnr, ite = continuous_nr(nE, y0.array) -# -# -# sys_df = pd.read_excel('instances/4node3pipe_bench.xlsx', -# sheet_name=None, -# engine='openpyxl', -# header=None -# ) -# -# -# def test_nr_method(): -# for var_name in ['Ts', 'Tr', 'm', 'mq', 'phi']: -# # find nonzero elements -# idx_nonzero = np.nonzero(sol.y[var_name]) -# assert max(abs((sol.y[var_name][idx_nonzero] - np.asarray(sys_df[var_name])[idx_nonzero].reshape(-1, ))) / -# np.asarray(sys_df[var_name])[idx_nonzero].reshape(-1, )) <= 1e-8 -# -# # def test_cnr_method(): -# # for var_name in ['Ts', 'Tr', 'm', 'mq', 'phi']: -# # # find nonzero elements -# # idx_nonzero = np.nonzero(y_cnr[var_name]) -# # assert max(abs((y_cnr[var_name][idx_nonzero] - np.asarray(sys_df[var_name])[idx_nonzero].reshape(-1, )) / -# # np.asarray(sys_df[var_name])[idx_nonzero].reshape(-1, ))) <= 1e-8 diff --git a/tests/test_dhspf_change_sign.py b/tests/test_dhspf_change_sign.py deleted file mode 100644 index ba75dbf..0000000 --- a/tests/test_dhspf_change_sign.py +++ /dev/null @@ -1,160 +0,0 @@ -# import pandas as pd -# import numpy as np -# from functools import partial -# -# from Solverz import Eqn, AE, nr_method, as_Vars, iVar, Param, continuous_nr, Para, idx, Abs, exp, Mat_Mul, \ -# made_numerical -# -# # %% initialize variables and params -# sys_df = pd.read_excel('instances/4node3pipe_change_sign.xlsx', -# sheet_name=None, -# engine='openpyxl', -# index_col=0 -# ) -# -# -# def derive_incidence_matrix(f_node, t_node): -# V = np.zeros((max(max(f_node), max(t_node)) + 1, len(f_node))) -# for pipe in range(len(f_node)): -# V[f_node[pipe], pipe] = -1 -# V[t_node[pipe], pipe] = 1 -# return V -# -# -# def derive_v_plus(v: np.ndarray) -> np.ndarray: -# return (v + np.abs(v)) / 2 -# -# -# def derive_v_minus(v: np.ndarray) -> np.ndarray: -# return (v - np.abs(v)) / 2 -# -# -# def v_p_reverse_pipe(V0: np.ndarray, m0: np.ndarray, m: np.ndarray) -> np.ndarray: -# """ -# Trigger function of v_p -# if some element of m changes its sign, then related column of V0 changes its sign -# :param V0: -# :param m0: initial mass flow rate -# :param m: -# :return: $V_0*(I-diag(sign(m0)-sign(m)))$ -# """ -# # pipe i changes sign then m_sign[i]=2 -# m_sign = np.abs(np.sign(m0) - np.sign(m)) -# return derive_v_plus(V0 @ (np.eye(m.shape[0]) - np.diagflat(m_sign))) -# -# -# def v_m_reverse_pipe(V0: np.ndarray, m0: np.ndarray, m: np.ndarray) -> np.ndarray: -# # pipe i changes sign then m_sign[i]=2 -# m_sign = np.abs(np.sign(m0) - np.sign(m)) -# return derive_v_minus(V0 @ (np.eye(m.shape[0]) - np.diagflat(np.abs(m_sign)))) -# -# -# def f_node_calculator(f_node0: np.ndarray, t_node0: np.ndarray, m0: np.ndarray, m: np.ndarray) -> np.ndarray: -# # pipe i changes sign then m_sign[i]=2 -# m_sign = np.abs(np.sign(m0) - np.sign(m)) -# flag_change_sign = m_sign > 0 -# return np.where(flag_change_sign.reshape((-1,)), t_node0, f_node0) -# -# -# def t_node_calculator(f_node0: np.ndarray, t_node0: np.ndarray, m0: np.ndarray, m: np.ndarray) -> np.ndarray: -# # pipe i changes sign then m_sign[i]=2 -# m_sign = np.abs(np.sign(m0) - np.sign(m)) -# flag_change_sign = m_sign > 0 -# return np.where(flag_change_sign.reshape((-1,)), f_node0, t_node0) -# -# -# # %% md -# # Declare equations and parameters -# # %% -# Node = sys_df['Node'] -# i = idx(name='i', value=np.asarray(Node['type'][Node['type'] == 3].index)) # intermediate nodes -# l = idx(name='l', value=np.asarray(Node['type'][Node['type'] == 2].index)) # load nodes -# s = idx(name='s', value=np.asarray(Node['type'][Node['type'] == 1].index)) # non-balance source nodes -# r = idx(name='r', value=np.asarray(Node['type'][Node['type'] == 0].index)) # balance source nodes -# sl = idx(name='sl', value=np.concatenate((s.value, l.value))) -# rs = idx(name='rs', value=np.concatenate((r.value, s.value))) -# rsl = idx(name='rsl', value=np.concatenate((r.value, s.value, l.value))) -# li = idx(name='li', value=np.concatenate((l.value, i.value))) -# rsi = idx(name='rsi', value=np.concatenate((r.value, s.value, i.value))) -# -# mL = Para('mL', dim=2, value=np.asarray(sys_df['m_L'])) -# K = Para(name='K', value=np.asarray(sys_df['K']).reshape(-1, )) -# phi_set = Para(name='phi_set', value=np.asarray(sys_df['phi_set']).reshape(-1, )) -# -# Cp = Para(name='Cp', value=np.asarray(sys_df['Cp']).reshape(-1, )) -# L = Para(name='L', value=np.asarray(sys_df['L']).reshape(-1, )) -# coeff_lambda = Para(name='coeff_lambda', value=np.asarray(sys_df['coeff_lambda']).reshape(-1, )) -# Ta = Para(name='Ta', value=np.asarray(sys_df['Ta']).reshape(-1, )) -# f_node = sys_df['Pipe']['f_node'] -# t_node = sys_df['Pipe']['t_node'] -# V = Para(name='V', dim=2, value=derive_incidence_matrix(f_node, t_node)) -# Vp = Para(name='Vp', dim=2, value=derive_v_plus(V.value)) -# Vm = Para(name='Vm', dim=2, value=derive_v_minus(V.value)) -# f_node = idx(name='f_node', value=np.asarray(f_node).reshape(-1, )) # intermediate nodes -# t_node = idx(name='t_node', value=np.asarray(t_node).reshape(-1, )) # load nodes -# -# m = iVar(name='m', value=np.asarray(sys_df['var']['m'])) -# mq = iVar(name='mq', value=np.asarray(sys_df['var']['mq'])) -# Ts = iVar(name='Ts', value=np.asarray(sys_df['var']['Ts'])) -# Tr = iVar(name='Tr', value=np.asarray(sys_df['var']['Tr'])) -# Touts = iVar(name='Touts', value=np.asarray(sys_df['var']['Touts'])) -# Toutr = iVar(name='Toutr', value=np.asarray(sys_df['var']['Toutr'])) -# phi = iVar(name='phi', value=np.asarray(sys_df['var']['phi'])) -# -# E1 = Eqn(name='E1', eqn=(Ts[f_node] - Ta) * exp(-coeff_lambda * L / (Cp * Abs(m))) + Ta - Touts) -# E3 = Eqn(name='E3', eqn=(Tr[t_node] - Ta) * exp(-coeff_lambda * L / (Cp * Abs(m))) + Ta - Toutr) -# E5 = Eqn(name='E5', eqn=Mat_Mul(V[rs, :], m) + mq[rs]) -# E6 = Eqn(name='E6', eqn=Mat_Mul(V[l, :], m) - mq[l]) -# E7 = Eqn(name='E7', eqn=Mat_Mul(V[i, :], m)) -# E8 = Eqn(name='E8', eqn=Mat_Mul(mL, K * Abs(m) * m)) -# E9 = Eqn(name='E9', eqn=Ts[li] * Mat_Mul(Vp[li, :], Abs(m)) - Mat_Mul(Vp[li, :], Touts * Abs(m))) -# E10 = Eqn(name='E10', eqn=Tr[rsi] * Mat_Mul(Vm[rsi, :], Abs(m)) - Mat_Mul(Vm[rsi, :], Toutr * Abs(m))) -# E11 = Eqn(name='E11', eqn=phi[rsl] - 4182 * mq[rsl] * (Ts[rsl] - Tr[rsl])) -# E12 = Eqn(name='E12', eqn=Ts[rs] - 100) -# E13 = Eqn(name='E13', eqn=Tr[l] - 50) -# E14 = Eqn(name='E14', eqn=phi[sl] - phi_set) -# E15 = Eqn(name='E15', eqn=phi[i]) -# E16 = Eqn(name='E16', eqn=mq[i]) -# E = AE(name='Pipe Equations', eqn=[E1, E3, E5, E6, E7, E8, E9, E10, E11, E12, E13, E14, E15, E16]) -# -# E.param_initializer('Vp', param=Param('Vp', value=Vp.value, triggerable=True, trigger_var='m', -# trigger_fun=partial(v_p_reverse_pipe, V.value, m.value), dim=2)) -# -# E.param_initializer('Vm', param=Param('Vm', value=Vm.value, triggerable=True, trigger_var='m', -# trigger_fun=partial(v_m_reverse_pipe, V.value, m.value), dim=2)) -# -# E.param_initializer('f_node', param=Param('f_node', value=f_node.value, triggerable=True, trigger_var='m', -# trigger_fun=partial(f_node_calculator, -# f_node.value, t_node.value, m.value), dtype=int)) -# -# E.param_initializer('t_node', param=Param('t_node', value=t_node.value, triggerable=True, trigger_var='m', -# trigger_fun=partial(t_node_calculator, -# f_node.value, t_node.value, m.value), dtype=int)) -# -# y0 = as_Vars([m, mq, Ts, Tr, Touts, Toutr, phi]) -# nE, code = made_numerical(E, y0, output_code=True) -# -# sol_nr = nr_method(nE, y0) -# # y_cnr, ite = continuous_nr(nE, y0.array) -# -# -# sys_df = pd.read_excel('instances/4node3pipe_change_sign_bench.xlsx', -# sheet_name=None, -# engine='openpyxl', -# header=None) -# -# -# def test_nr_method(): -# for var_name in ['Ts', 'Tr', 'm', 'mq', 'phi']: -# # find nonzero elements -# idx_nonzero = np.nonzero(sol_nr.y[var_name]) -# assert max(abs((sol_nr.y[var_name][idx_nonzero] - np.asarray(sys_df[var_name])[idx_nonzero].reshape(-1, ))) / -# np.asarray(sys_df[var_name])[idx_nonzero].reshape(-1, )) <= 1e-8 -# -# -# # def test_cnr_method(): -# # for var_name in ['Ts', 'Tr', 'm', 'mq', 'phi']: -# # # find nonzero elements -# # idx_nonzero = np.nonzero(y_cnr[var_name]) -# # assert max(abs((y_cnr[var_name][idx_nonzero] - np.asarray(sys_df[var_name])[idx_nonzero].reshape(-1, )) / -# # np.asarray(sys_df[var_name])[idx_nonzero].reshape(-1, ))) <= 1e-8 diff --git a/tests/test_dyn_gas_flow_single_pipe.py b/tests/test_dyn_gas_flow_single_pipe.py deleted file mode 100644 index 9253ee1..0000000 --- a/tests/test_dyn_gas_flow_single_pipe.py +++ /dev/null @@ -1,64 +0,0 @@ -import numpy as np -import pandas as pd - -from Solverz import idx, iVar, Para, as_Vars, HyperbolicPde, fdae_solver, Eqn, Param, IdxParam, FDAE, made_numerical, Opt - -results = pd.read_excel('instances/dynamic_gas_flow_single_pipe.xlsx', - sheet_name=None, - engine='openpyxl' - ) - -Pie = iVar('Pie', value=np.linspace(9, 8.728, 11) * 1e6) -q = iVar('q', value=7.3125 * np.ones((11,))) -va = Para('va', value=340) -D = Para('D', value=0.5) -S = Para('S', value=np.pi * (D.value / 2) ** 2) -lam = Para('lam', 0.3) -pb = Para('pb', 9e6) -qb = Para('qb', 7.3125) -M = idx('M') -L = 3000 - -gas_pde1 = HyperbolicPde('mass conservation', - diff_var=Pie, - flux=va ** 2 / S * q, - two_dim_var=[Pie, q]) -gas_pde2 = HyperbolicPde('momentum conservation', - diff_var=q, - flux=S * Pie, - source=-lam * va ** 2 * q ** 2 / (2 * D * S * Pie), - two_dim_var=[Pie, q]) - -ae1 = gas_pde1.finite_difference() -ae2 = gas_pde2.finite_difference() -ae3 = Eqn(name='boundary condition Pie', eqn=Pie[0] - pb) -ae4 = Eqn(name='boundary condition q', eqn=q[M] - qb) - -gas_FDE = FDAE([ae1, ae2, ae3, ae4], name='gas FDE', nstep=1) -u0 = as_Vars([Pie, q]) -gas_FDE.update_param(u0.derive_alias('_tag_0')) - -gas_FDE.param_initializer('M', IdxParam('M', value=np.array(10))) -gas_FDE.param_initializer('dx', Param('dx', value=np.array(300))) -gas_FDE.param_initializer('dt', Param('dt', value=5)) - -ngas = made_numerical(gas_FDE, u0) -sol1 = fdae_solver(ngas, [0, 3600], u0, Opt(step_size=5, ite_tol=1e-3)) -ngas_sparse, code = made_numerical(gas_FDE, u0, sparse=True, output_code=True) - -sol2 = fdae_solver(ngas_sparse, [0, 3600], u0, Opt(step_size=5, ite_tol=1e-3)) - - -def test_fde_solver(): - pout = np.asarray(results['results']['pout']) - qin = np.asarray(results['results']['qin']) - - # if matrix container is cvxopt.spmatrix - # assert np.mean(abs(pout - u['Pie'][:, -1] / 1e6)) < 5e-6 - # assert np.mean(abs(qin - u['q'][:, 0])) < 1.2e-4 - - # if matrix container is scipy.sparse.csc_array - assert np.mean(abs(pout - sol1.Y['Pie'][:, -1] / 1e6)) < 1e-10 - assert np.mean(abs(qin - sol1.Y['q'][:, 0])) < 1e-7 - assert np.mean(abs(pout - sol2.Y['Pie'][:, -1] / 1e6)) < 1e-10 - assert np.mean(abs(qin - sol2.Y['q'][:, 0])) < 1e-7 diff --git a/tests/test_gasflow.py b/tests/test_gasflow.py deleted file mode 100644 index 9b84577..0000000 --- a/tests/test_gasflow.py +++ /dev/null @@ -1,33 +0,0 @@ -import numpy as np - -from Solverz import idx, iVar, Para, Sign, as_Vars, nr_method, Eqn, AE, Mat_Mul, made_numerical - -k = idx('k', value=[0, 1, 2]) -i = idx('i', value=[0, 0, 3]) -j = idx('j', value=[1, 2, 2]) -m = idx('m', value=[1, 2, 3]) - -Pi = iVar('Pi', value=[50, 49, 45, 1.3 * 49]) # 50, 40.8160, 49.7827, 1.3 * 40.8160 -fin = iVar('fin', value=[34.6406, 10.8650, 23.7756, 0]) # 29.8308, 4.8098, 18.9658 -finset = Para('finset', value=[10.8650, 23.7756, 0]) -f = iVar('f', value=[10, 5, 25, 25]) -c = Para('c', value=[1.0329225961928894, 1.0329267609699861, 1.032931789457253, 1]) -A = Para('A', dim=2, value=[[-1, -1, 0, 0], [1, 0, 0, -1], [0, 1, 1, 0], [0, 0, -1, 1]]) - -eqn1 = Eqn(name='mass flow continuity', eqn=finset - Mat_Mul(A[m, :], f)) -eqn2 = Eqn(name='mass flow1', - eqn=f[k] - c[k] * Sign(Pi[i] - Pi[j]) * (Sign(Pi[i] - Pi[j]) * (Pi[i] ** 2 - Pi[j] ** 2)) ** (1 / 2)) -eqn3 = Eqn(name='pressure', eqn=Pi[0] - 50) -eqn4 = Eqn(name='pressure1', eqn=Pi[3] - 1.3 * Pi[1]) - -gas_flow = AE(eqn=[eqn1, eqn2, eqn3, eqn4]) -y0 = as_Vars([f, Pi]) - -ngas_flow = made_numerical(gas_flow, y0) -sol = nr_method(ngas_flow, y0) - - -def test_nr_method(): - bench = np.array([29.830799999999996, 4.809800000000001, 18.965799999999998, - 18.965799999999998, 50.0, 40.816, 49.78269999999999, 53.06080000000001]) - assert np.max(np.abs((sol.y - bench) / bench)) <= 1e-8 diff --git a/tests/test_m3b9.py b/tests/test_m3b9.py deleted file mode 100644 index 0b6a1aa..0000000 --- a/tests/test_m3b9.py +++ /dev/null @@ -1,101 +0,0 @@ -import numpy as np -import pandas as pd - -from Solverz import DAE, Eqn, Ode, iVar, Para, idx, sin, cos, Rodas, as_Vars, Mat_Mul, \ - implicit_trapezoid, Opt, TimeSeriesParam, made_numerical - -omega = iVar(name='omega') -delta = iVar(name='delta') -Ux = iVar(name='Ux') -Uy = iVar(name='Uy') -Ixg = iVar(name='Ixg') -Iyg = iVar(name='Iyg') -g = idx('g', value=[0, 1, 2]) -ng = idx('ng', value=[3, 4, 5, 6, 7, 8]) -Pm = Para(name='Pm') -G = Para(name='G', dim=2) -B = Para(name='B', dim=2) -D = Para(name='D') -Tj = Para(name='Tj') -ra = Para(name='ra') -omega_b = Para(name='omega_b') -Edp = Para(name='Edp') -Eqp = Para(name='Eqp') -Xdp = Para(name='Xdp') -Xqp = Para(name='Xqp') - -rotator_eqn = Ode(name='rotator speed', - f=(Pm - (Ux[g] * Ixg + Uy[g] * Iyg + (Ixg ** 2 + Iyg ** 2) * ra) - D * (omega - 1)) / Tj, - diff_var=omega) -delta_eqn = Ode(name='delta', f=(omega - 1) * omega_b, diff_var=delta) -Ed_prime = Eqn(name='Ed_prime', eqn=Edp - sin(delta) * (Ux[g] + ra * Ixg - Xqp * Iyg) + cos(delta) * ( - Uy[g] + ra * Iyg + Xqp * Ixg)) -Eq_prime = Eqn(name='Eq_prime', eqn=Eqp - cos(delta) * (Ux[g] + ra * Ixg - Xdp * Iyg) - sin(delta) * ( - Uy[g] + ra * Iyg + Xdp * Ixg)) -Ix_inject = Eqn(name='Ixg_inject', eqn=Ixg - (Mat_Mul(G[g, :], Ux) - Mat_Mul(B[g, :], Uy))) -Iy_inject = Eqn(name='Iyg_inject', eqn=Iyg - (Mat_Mul(G[g, :], Uy) + Mat_Mul(B[g, :], Ux))) -Ixng_inject = Eqn(name='Ixng_inject', eqn=Mat_Mul(G[ng, :], Ux) - Mat_Mul(B[ng, :], Uy)) -Iyng_inject = Eqn(name='Iyng_inject', eqn=Mat_Mul(G[ng, :], Uy) + Mat_Mul(B[ng, :], Ux)) - -dt = np.array(0.002) -T = np.array(10) - -m3b9 = DAE([delta_eqn, rotator_eqn, Ed_prime, Eq_prime, Ix_inject, Iy_inject, Ixng_inject, Iyng_inject], - name='m3b9') -m3b9.update_param('Pm', [0.7164, 1.6300, 0.8500]) -m3b9.update_param('ra', [0.0000, 0.0000, 0.0000]) -m3b9.update_param('D', [10, 10, 10]) -m3b9.update_param('Tj', [47.2800, 12.8000, 6.0200]) -m3b9.update_param('omega_b', [376.991118430775]) -m3b9.update_param('g', [0, 1, 2]) -m3b9.update_param('ng', [3, 4, 5, 6, 7, 8]) -m3b9.update_param('Edp', [0.0000, 0.0000, 0.0000]) -m3b9.update_param('Eqp', [1.05636632091501, 0.788156757672709, 0.767859471854610]) -m3b9.update_param('Xdp', [0.0608, 0.1198, 0.1813]) -m3b9.update_param('Xqp', [0.0969, 0.8645, 1.2578]) -df = pd.read_excel('instances/test_m3b9.xlsx', - sheet_name=None, - engine='openpyxl', - header=None - ) -Gvalue = np.asarray(df['G']) -m3b9.param_initializer('G', TimeSeriesParam(name='G', - v_series=[Gvalue[6, 6], 10000, 10000, Gvalue[6, 6], Gvalue[6, 6]], - time_series=[0, 0.002, 0.03, 0.032, T], - value=Gvalue, - index=(6, 6), - dim=2)) -m3b9.update_param('B', np.asarray(df['B'])) - -delta = iVar('delta', [0.0625815077879868, 1.06638275203221, 0.944865048677501]) -omega = iVar('omega', [1, 1, 1]) -Ixg = iVar('Ixg', [0.688836021737262, 1.57988988391346, 0.817891311823357]) -Iyg = iVar('Iyg', [-0.260077644814056, 0.192406178191528, 0.173047791590276]) -Ux = iVar('Ux', - [1.04000110267534, 1.01157932564567, 1.02160343921907, - 1.02502063033405, 0.993215117729926, 1.01056073782038, - 1.02360471178264, 1.01579907336413, 1.03174403980626]) -Uy = iVar('Uy', - [9.38510394478286e-07, 0.165293826097057, 0.0833635520284917, - -0.0396760163416718, -0.0692587531054159, -0.0651191654677445, - 0.0665507083524658, 0.0129050646926083, 0.0354351211556429]) - -y0 = as_Vars([delta, omega, Ixg, Iyg, Ux, Uy]) - -m3b9_dae, code = made_numerical(m3b9, y0, sparse=False, output_code=True) -sol_rodas = Rodas(m3b9_dae, - np.linspace(0, 10, 5001).reshape(-1, ), - y0, - Opt(hinit=1e-5)) - -sol_trape = implicit_trapezoid(m3b9_dae, - [0, 10], - y0, - Opt(step_size=dt)) - - -def test_m3b9(): - assert np.abs(np.asarray(df['omega']) - sol_rodas.Y['omega']).max() <= 5e-5 - assert np.abs(np.asarray(df['omega']) - sol_trape.Y['omega']).max() <= 2.48e-4 - assert np.abs(np.asarray(df['delta']) - sol_rodas.Y['delta']).max() <= 9.6e-4 - assert np.abs(np.asarray(df['delta']) - sol_trape.Y['delta']).max() <= 7.8e-2