zoomy-core 0.1.1__py3-none-any.whl → 0.1.3__py3-none-any.whl

zoomy_core/transformation/to_amrex.py

@@ -0,0 +1,238 @@
+import os
+from sympy import MatrixSymbol, fraction, cancel, Matrix, symbols, radsimp, powsimp
+import sympy as sp
+from copy import deepcopy
+
+from library.zoomy_core.misc.misc import Zstruct
+from library.zoomy_core.model.sympy2c import create_module
+from library.zoomy_core.transformation.helpers import regularize_denominator, substitute_sympy_attributes_with_symbol_matrix
+
+import sympy as sp
+from sympy.printing.cxx import CXX11CodePrinter
+import re
+import textwrap
+
+
+class AmrexPrinter(CXX11CodePrinter):
+    """
+    After the normal C++ printer has done its job, replace every
+    'std::foo(' with 'amrex::Math::foo('  – except if foo is listed
+    in 'custom_map'.  No other overrides are necessary.
+    """
+    def __init__(self, model, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.map_Q = {k: f"Q({i})" for i, k in enumerate(model.variables.values())}
+        self.map_Qaux = {k: f"Qaux({i})" for i, k in enumerate(model.aux_variables.values())}
+        self.map_param = {k: str(float(model.parameter_values[i])) for i, k in enumerate(model.parameters.values())}
+        self.map_normal = {k:  f"normal({i})" for i, k in enumerate(model.normal.values())}
+        self.map_position = {k:  f"X({i})" for i, k in enumerate(model.position.values())}
+
+
+        self._custom_map = set({})
+        # names that should *keep* their std:: prefix
+
+        # pre-compile regex  std::something(
+        self._std_regex = re.compile(r'std::([A-Za-z_]\w*)')
+        
+    def _print_Symbol(self, s):
+        for map in [self.map_Q, self.map_Qaux, self.map_param, self.map_normal, self.map_position]:
+            if s in map:
+                return map[s]
+        return super()._print_Symbol(s)
+    
+    def _print_Pow(self, expr):
+        """
+        Print a SymPy Power.
+
+        * integer exponent  -> amrex::Math::powi<EXP>(base)
+        * otherwise         -> amrex::Math::pow (run-time exponent)
+        """
+        base, exp = expr.as_base_exp()
+
+        # integer exponent ------------------------------------------------
+        if exp.is_Integer:
+            n = int(exp)
+
+            # 0, 1 and negative exponents inlined
+            if n == 0:
+                return "1.0"
+            if n == 1:
+                return self._print(base)
+            if n < 0:
+                # negative integer: 1 / powi<-n>(base)
+                return (f"(1.0 / amrex::Math::powi<{abs(n)}>("
+                        f"{self._print(base)}))")
+
+            # positive integer
+            return f"amrex::Math::powi<{n}>({self._print(base)})"
+
+        # non-integer exponent -------------------------------------------
+        return (f"std::pow("
+                f"{self._print(base)}, {self._print(exp)})")
+
+    # the only method we override
+    def doprint(self, expr, **settings):
+        code = super().doprint(expr, **settings)
+
+        # callback that the regex will call for every match
+        def _repl(match):
+            fname = match.group(1)
+            if fname in self._custom_map:
+                return self._custom_map[fname]
+            else:
+                return f'std::{fname}'
+
+        # apply the replacement to the whole code string
+        return self._std_regex.sub(_repl, code) 
+
+    def convert_expression_body(self, expr, target='res'):
+
+        tmp_sym   = sp.numbered_symbols('t') 
+        temps, simplified = sp.cse(expr, symbols=tmp_sym)  
+        lines = []
+        for lhs, rhs in temps:
+            lines.append(f"amrex::Real {self.doprint(lhs)} = {self.doprint(rhs)};")
+
+        for i in range(expr.rows):
+            for j in range(expr.cols):
+                lines.append(f"{target}({i},{j}) = {self.doprint(simplified[0][i, j])};")
+
+        body = '\n        '.join(lines)
+        return body
+    
+    def createSmallMatrix(self, rows, cols):
+        return f"amrex::SmallMatrix<amrex::Real,{rows},{cols}>"
+    
+    def create_file_header(self, n_dof_q, n_dof_qaux, dim):
+        header = textwrap.dedent(f"""
+        #pragma once
+        #include <AMReX_Array4.H>
+        #include <AMReX_Vector.H>
+        
+        class Model {{
+        public:
+            static constexpr int n_dof_q    = {n_dof_q};
+            static constexpr int n_dof_qaux = {n_dof_qaux};
+            static constexpr int dimension  = {dim};
+        """)
+        return header
+    
+    def create_file_footer(self):
+        return  """
+};
+                """
+    
+    def create_function(self, name, expr, n_dof_q, n_dof_qaux, target='res'):
+        if type(expr) is list:
+            dim = len(expr)
+            return [self.create_function(f"{name}_{dir}", expr[i], n_dof_q, n_dof_qaux) for i, dir in enumerate(['x', 'y', 'z'][:dim])]
+        res_shape = expr.shape
+        body = self.convert_expression_body((expr), target=target)
+        text = f"""
+    AMREX_GPU_HOST_DEVICE
+    AMREX_FORCE_INLINE
+    static {self.createSmallMatrix(*res_shape)}
+    {name} ( {self.createSmallMatrix(n_dof_q, 1)} const& Q,
+    {self.createSmallMatrix(n_dof_qaux, 1)} const& Qaux) noexcept
+    {{
+        auto {target} = {self.createSmallMatrix(*res_shape)}{{}};
+        {body}
+        return {target};
+    }}
+        """
+        return text
+
+    def create_function_normal(self, name, expr, n_dof_q, n_dof_qaux, dim, target='res'):
+        if type(expr) is list:
+            dim = len(expr)
+            return [self.create_function_normal(f"{name}_{dir}", expr[i], n_dof_q, n_dof_qaux, dim) for i, dir in enumerate(['x', 'y', 'z'][:dim])]
+        res_shape = expr.shape
+        body = self.convert_expression_body(expr, target=target)
+        text = f"""
+    AMREX_GPU_HOST_DEVICE
+    AMREX_FORCE_INLINE
+    static {self.createSmallMatrix(*res_shape)}
+    {name} ( {self.createSmallMatrix(n_dof_q, 1)} const& Q,
+    {self.createSmallMatrix(n_dof_qaux, 1)} const& Qaux,
+    {self.createSmallMatrix(dim, 1)} const& normal) noexcept
+    {{
+        auto {target} = {self.createSmallMatrix(*res_shape)}{{}};
+        {body}
+        return {target};
+
+    }}
+        """
+        return text
+    
+    def create_function_interpolate(self, name, expr, n_dof_q, n_dof_qaux, target='res'):
+        res_shape = expr.shape
+        body = self.convert_expression_body(expr, target=target)
+        text = f"""
+    AMREX_GPU_HOST_DEVICE
+    AMREX_FORCE_INLINE
+    static {self.createSmallMatrix(*res_shape)}
+    {name} ( {self.createSmallMatrix(n_dof_q, 1)} const& Q,
+    {self.createSmallMatrix(n_dof_qaux, 1)} const& Qaux,
+    {self.createSmallMatrix(3, 1)} const& X) noexcept
+    {{
+        auto {target} = {self.createSmallMatrix(*res_shape)}{{}};
+        {body}
+        return {target};
+    }}
+
+        """
+        return text
+
+
+    def create_function_boundary(self, name, expr, n_dof_q, n_dof_qaux, dim, target='res'):
+        res_shape = expr.shape
+        body = self.convert_expression_body(expr, target=target)
+        text = f"""
+    AMREX_GPU_HOST_DEVICE
+    AMREX_FORCE_INLINE
+    static {self.createSmallMatrix(*res_shape)}
+    {name} ( {self.createSmallMatrix(n_dof_q, 1)} const& Q,
+    {self.createSmallMatrix(n_dof_qaux, 1)} const& Qaux,
+    {self.createSmallMatrix(dim, 1)} const& normal, 
+    {self.createSmallMatrix(3, 1)} const& position,
+    amrex::Real const& time,
+    amrex::Real const& dX) noexcept
+    {{
+        auto {target} = {self.createSmallMatrix(*res_shape)}{{}};
+        {body}
+        return {target};
+
+    }}
+        """
+        return text
+    
+    def create_model(self, model):
+        n_dof = model.n_variables
+        n_dof_qaux = model.n_aux_variables
+        dim =  model.dimension
+        module_functions = []
+        module_functions += self.create_function('flux', model.flux(), n_dof, n_dof_qaux)
+        module_functions += self.create_function('flux_jacobian', model.flux_jacobian(), n_dof, n_dof_qaux)
+        module_functions += self.create_function('nonconservative_matrix', model.nonconservative_matrix(), n_dof, n_dof_qaux)
+        module_functions += self.create_function('quasilinear_matrix', model.quasilinear_matrix(), n_dof, n_dof_qaux)
+        module_functions.append(self.create_function_normal('eigenvalues', model.eigenvalues(), n_dof, n_dof_qaux, dim))
+        module_functions.append(self.create_function('left_eigenvectors', model.left_eigenvectors(), n_dof, n_dof_qaux))
+        module_functions.append(self.create_function('right_eigenvectors', model.right_eigenvectors(), n_dof, n_dof_qaux))
+        module_functions.append(self.create_function('source', model.source(), n_dof, n_dof_qaux))
+        module_functions.append(self.create_function('residual', model.residual(), n_dof, n_dof_qaux))
+        module_functions.append(self.create_function('source_implicit', model.source_implicit(), n_dof, n_dof_qaux))
+        module_functions.append(self.create_function_interpolate('project_2d_to_3d', model.project_2d_to_3d(), n_dof, n_dof_qaux))
+        module_functions.append(self.create_function_boundary('boundary_conditions', model.boundary_conditions.get_boundary_condition_function(model.time, model.position, model.distance, model.variables, model.aux_variables, model.parameters, model.normal), n_dof, n_dof_qaux, dim))
+        full = self.create_file_header(n_dof, n_dof_qaux, dim) + '\n\n' + '\n\n'.join(module_functions) + self.create_file_footer()
+        return full
+    
+
+def write_code(model, settings):
+    printer = AmrexPrinter(model)
+    expr = printer.create_model(model)
+    main_dir = os.getenv("ZOOMY_DIR")
+    path = os.path.join(main_dir, settings.output.directory, ".amrex_interface")
+    os.makedirs(path, exist_ok=True)
+    path = os.path.join(path, "Model.h")
+    with open(path, 'w+') as f:
+        f.write(expr)

zoomy_core/transformation/to_c.py

@@ -0,0 +1,181 @@
+import os
+import re
+import textwrap
+import sympy as sp
+from sympy.printing.cxx import CXX11CodePrinter
+
+class CPrinter(CXX11CodePrinter):
+    """
+    Convert SymPy expressions to C code.
+    """
+
+    def __init__(self, model, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self.n_dof_q = model.n_variables
+        self.n_dof_qaux = model.n_aux_variables
+
+        # Map variable names to Q[i], Qaux[i]
+        self.map_Q = {k: f"Q[{i}]" for i, k in enumerate(model.variables.values())}
+        self.map_Qaux = {k: f"Qaux[{i}]" for i, k in enumerate(model.aux_variables.values())}
+        self.map_param = {k: str(float(model.parameter_values[i])) for i, k in enumerate(model.parameters.values())}
+        
+        self.map_normal = {k: f"n[{i}]" for i, k in enumerate(model.normal.values())}
+        self.map_position = {k: f"X[{i}]" for i, k in enumerate(model.position.values())}
+
+        self._std_regex = re.compile(r'std::([A-Za-z_]\w*)')
+
+    # --- Symbol printing --------------------------------------------------
+    def _print_Symbol(self, s):
+        for m in [self.map_Q, self.map_Qaux, self.map_param, self.map_normal, self.map_position]:
+            if s in m:
+                return m[s]
+        return super()._print_Symbol(s)
+
+    # --- Pow printing -----------------------------------------------------
+    def _print_Pow(self, expr):
+        base, exp = expr.as_base_exp()
+        if exp.is_Integer:
+            n = int(exp)
+            if n == 0:
+                return "1.0"
+            if n == 1:
+                return self._print(base)
+            if n < 0:
+                return f"(1.0 / std::pow({self._print(base)}, {abs(n)}))"
+            return f"std::pow({self._print(base)}, {n})"
+        return f"std::pow({self._print(base)}, {self._print(exp)})"
+
+    # --- Expression conversion --------------------------------------------
+    def convert_expression_body(self, expr, target='res'):
+        tmp_sym = sp.numbered_symbols('t')
+        temps, simplified = sp.cse(expr, symbols=tmp_sym)
+        lines = []
+        cols = expr.shape[1]
+        for lhs, rhs in temps:
+            lines.append(f"double {self.doprint(lhs)} = {self.doprint(rhs)};")
+        for i in range(expr.rows):
+            for j in range(expr.cols):
+                lines.append(f"{target}[{i * cols +j}] = {self.doprint(simplified[0][i, j])};")
+        return "\n        ".join(lines)
+
+    # --- Header / Footer --------------------------------------------------
+    def create_file_header(self, n_dof_q, n_dof_qaux, dim, list_sorted_function_names):
+        return textwrap.dedent(f"""\
+        #pragma once
+
+        static const int MODEL_n_dof_q    = {n_dof_q};
+        static const int MODEL_n_dof_qaux = {n_dof_qaux};
+        static const int MODEL_dimension  = {dim};
+        static const int MODEL_n_boundary_tags = {len(list_sorted_function_names)};
+        static const char* MODEL_map_boundary_tag_to_function_index[] = {{ {", ".join(f'"{item}"' for item in list_sorted_function_names)} }};
+        """)
+
+    def create_file_footer(self):
+        return "\n"
+
+    # --- Function generators ---------------------------------------------
+    def create_function(self, name, expr, n_dof_q, n_dof_qaux, target='res'):
+        if isinstance(expr, list):
+            dim = len(expr)
+            return [self.create_function(f"{name}_{d}", expr[i], n_dof_q, n_dof_qaux)
+                    for i, d in enumerate(['x', 'y', 'z'][:dim])]
+
+        rows, cols = expr.shape
+        body = self.convert_expression_body(expr, target)
+        return f"""
+inline void {name}(
+    const double* Q,
+    const double* Qaux,
+    double* res
+    )
+{{
+    {body}
+}}
+        """
+
+    def create_function_normal(self, name, expr, n_dof_q, n_dof_qaux, dim, target='res'):
+        if isinstance(expr, list):
+            return [self.create_function_normal(f"{name}_{d}", expr[i], n_dof_q, n_dof_qaux, dim)
+                    for i, d in enumerate(['x', 'y', 'z'][:dim])]
+
+        rows, cols = expr.shape
+        body = self.convert_expression_body(expr, target)
+        return f"""
+inline void {name}(
+    const double* Q,
+    const double* Qaux,
+    const double* n,
+    double* res)
+{{
+    {body}
+}}
+        """
+
+    def create_function_interpolate(self, name, expr, n_dof_q, n_dof_qaux, target='res'):
+        rows, cols = expr.shape
+        body = self.convert_expression_body(expr, target)
+        return f"""
+inline double* {name}(
+    const double* Q,
+    const double* Qaux,
+    const double* X,
+    double* res)
+{{
+    {body}
+}}
+        """
+
+    def create_function_boundary(self, name, expr, n_dof_q, n_dof_qaux, dim, target='res'):
+        rows, cols = expr.shape
+        body = self.convert_expression_body(expr, target)
+        return f"""
+inline double* {name}(
+    const double* Q,
+    const double* Qaux,
+    const double* n,
+    const double* X,
+    const double time,
+    const double dX, 
+    double* res)
+{{
+    {body}
+}}
+        """
+
+    # --- Full model generation ------------------------------------------
+    def create_model(self, model):
+        n_dof = model.n_variables
+        n_dof_qaux = model.n_aux_variables
+        dim = model.dimension
+        funcs = []
+        funcs += self.create_function('flux', model.flux(), n_dof, n_dof_qaux)
+        funcs += self.create_function('flux_jacobian', model.flux_jacobian(), n_dof, n_dof_qaux)
+        funcs += self.create_function('nonconservative_matrix', model.nonconservative_matrix(), n_dof, n_dof_qaux)
+        funcs += self.create_function('quasilinear_matrix', model.quasilinear_matrix(), n_dof, n_dof_qaux)
+        funcs.append(self.create_function_normal('eigenvalues', model.eigenvalues(), n_dof, n_dof_qaux, dim))
+        funcs.append(self.create_function('left_eigenvectors', model.left_eigenvectors(), n_dof, n_dof_qaux))
+        funcs.append(self.create_function('right_eigenvectors', model.right_eigenvectors(), n_dof, n_dof_qaux))
+        funcs.append(self.create_function('source', model.source(), n_dof, n_dof_qaux))
+        funcs.append(self.create_function('residual', model.residual(), n_dof, n_dof_qaux))
+        funcs.append(self.create_function('source_implicit', model.source_implicit(), n_dof, n_dof_qaux))
+        funcs.append(self.create_function_interpolate('interpolate', model.project_2d_to_3d(), n_dof, n_dof_qaux))
+        funcs.append(self.create_function_boundary(
+            'boundary_conditions',
+            model.boundary_conditions.get_boundary_condition_function(
+                model.time, model.position, model.distance,
+                model.variables, model.aux_variables, model.parameters, model.normal),
+            n_dof, n_dof_qaux, dim))
+
+        return self.create_file_header(n_dof, n_dof_qaux, dim, model.boundary_conditions.list_sorted_function_names) + "\n".join(funcs) + self.create_file_footer()
+
+
+def write_code(model, settings):
+    printer = CPrinter(model)
+    code = printer.create_model(model)
+    main_dir = os.getenv("ZOOMY_DIR")
+    path = os.path.join(main_dir, settings.output.directory, ".c_interface")
+    os.makedirs(path, exist_ok=True)
+    file_path = os.path.join(path, "Model.H")
+    with open(file_path, "w+") as f:
+        f.write(code)

zoomy_core/transformation/to_jax.py

@@ -0,0 +1,14 @@
+from attrs import define
+try:
+    import jax.numpy as jnp
+except ImportError:
+    import numpy as jnp
+
+from library.zoomy_core.transformation.to_numpy import NumpyRuntimeModel
+
+
+
+@define(kw_only=False, slots=True, frozen=True)
+class JaxRuntimeModel(NumpyRuntimeModel):
+    module = {'ones_like': jnp.ones_like, 'zeros_like': jnp.zeros_like, 'array': jnp.array, 'squeeze': jnp.squeeze}
+    printer="jax"

zoomy_core/transformation/to_numpy.py

@@ -0,0 +1,115 @@
+from typing import Callable
+import numpy as np
+from attrs import define, field
+
+from library.zoomy_core.misc.custom_types import FArray
+
+
+@define(kw_only=False, slots=True, frozen=True)
+class NumpyRuntimeModel:
+    """Runtime model generated from a symbolic Model."""
+
+    # --- User-provided ---
+    model: object = field()
+
+    # --- Automatically derived ---
+    name: str = field(init=False)
+    dimension: int = field(init=False)
+    n_variables: int = field(init=False)
+    n_aux_variables: int = field(init=False)
+    n_parameters: int = field(init=False)
+    parameters: FArray = field(init=False)
+    flux: Callable = field(init=False)
+    dflux: Callable = field(init=False)
+    source: Callable = field(init=False)
+    source_jacobian_wrt_variables: Callable = field(init=False)
+    source_jacobian_wrt_aux_variables: Callable = field(init=False)
+    nonconservative_matrix: Callable = field(init=False)
+    quasilinear_matrix: Callable = field(init=False)
+    eigenvalues: Callable = field(init=False)
+    residual: Callable = field(init=False)
+    project_2d_to_3d: Callable = field(init=False)
+    project_3d_to_2d: Callable = field(init=False)
+    boundary_conditions: Callable = field(default=None, init=False)
+
+    left_eigenvectors: Callable = field(default=None, init=False)
+    right_eigenvectors: Callable = field(default=None, init=False)
+
+    # --- Constants ---
+    module = {
+        "ones_like": np.ones_like,
+        "zeros_like": np.zeros_like,
+        "array": np.array,
+        "squeeze": np.squeeze,
+    }
+    printer = "numpy"
+    
+
+    # ------------------------------------------------------------------
+
+    def __attrs_post_init__(self):
+        model = self.model
+
+        
+        flux=model._flux.lambdify(modules=[self.module, self.printer])
+        dflux=model._dflux.lambdify(modules=[self.module, self.printer])
+        nonconservative_matrix=model._nonconservative_matrix.lambdify(
+            modules=[self.module, self.printer]
+        )
+        quasilinear_matrix=model._quasilinear_matrix.lambdify(
+            modules=[self.module, self.printer]
+        )
+        eigenvalues=model._eigenvalues.lambdify(
+            modules=[self.module, self.printer]
+        )
+        left_eigenvectors=model._left_eigenvectors.lambdify(
+            modules=[self.module, self.printer]
+        )
+        right_eigenvectors=model._right_eigenvectors.lambdify(
+            modules=[self.module, self.printer]
+        )
+        source=model._source.lambdify(modules=[self.module, self.printer])
+        source_jacobian_wrt_variables=model._source_jacobian_wrt_variables.lambdify(
+            modules=[self.module, self.printer]
+        )
+        source_jacobian_wrt_aux_variables=model._source_jacobian_wrt_aux_variables.lambdify(
+            modules=[self.module, self.printer]
+        )
+        residual=model._residual.lambdify(modules=[self.module, self.printer])
+        project_2d_to_3d=model._project_2d_to_3d.lambdify(
+            modules=[self.module, self.printer]
+        )
+        project_3d_to_2d=model._project_3d_to_2d.lambdify(
+            modules=[self.module, self.printer]
+        )
+        bcs = model._boundary_conditions.lambdify(modules=[self.module, self.printer])
+
+
+        # --- Assign frozen fields -------------------------------------
+        object.__setattr__(self, "name", model.name)
+        object.__setattr__(self, "dimension", model.dimension)
+        object.__setattr__(self, "n_variables", model.n_variables)
+        object.__setattr__(self, "n_aux_variables", model.n_aux_variables)
+        object.__setattr__(self, "n_parameters", model.n_parameters)
+        object.__setattr__(self, "parameters", model.parameter_values)
+
+        object.__setattr__(self, "flux", flux)
+        object.__setattr__(self, "dflux", dflux)
+        object.__setattr__(self, "source", source)
+        object.__setattr__(
+            self, "source_jacobian_wrt_variables", source_jacobian_wrt_variables
+        )
+        object.__setattr__(
+            self,
+            "source_jacobian_wrt_aux_variables",
+            source_jacobian_wrt_aux_variables,
+        )
+        object.__setattr__(self, "nonconservative_matrix", nonconservative_matrix)
+        object.__setattr__(self, "quasilinear_matrix", quasilinear_matrix)
+        object.__setattr__(self, "eigenvalues", eigenvalues)
+        object.__setattr__(self, "left_eigenvectors", left_eigenvectors)
+        object.__setattr__(self, "right_eigenvectors", right_eigenvectors)
+        object.__setattr__(self, "residual", residual)
+        object.__setattr__(self, "project_2d_to_3d", project_2d_to_3d)
+        object.__setattr__(self, "project_3d_to_2d", project_3d_to_2d)
+        object.__setattr__(self, "boundary_conditions", bcs)