Skip to content

C++ Header Code Generation (mkpp.codegen)

mkpp.codegen

MKPP Code Generation Orchestrator.

This module provides the top-level generate_headers function that emits Kokkos C++ solver headers from a parsed mechanism definition. It delegates to the Jinja2 template engine for C++ emission and to focused submodules for tableau definitions and expression formatting.

Public API (backward-compatible): - generate_headers - generate_host_api_headers - SOLVER_COEFFICIENTS - RosenbrockTableau - format_eqn - get_A, get_C

RosenbrockTableau dataclass

Immutable coefficient tableau for a Rosenbrock solver.

Source code in src/mkpp/rosenbrock.py 
11
12
13
14
15
16
17
18
19
20
21
22
23
24
@dataclass(frozen=True)
class RosenbrockTableau:
    """Immutable coefficient tableau for a Rosenbrock solver."""

    name: str
    stages: int
    A: list[float]  # Strictly lower-triangular, row-wise: A(2,1), A(3,1), A(3,2), ...
    C: list[float]  # Same storage as A
    M: list[float]  # Solution update weights, length = stages
    E: list[float]  # Error estimate weights, length = stages
    Alpha: list[float]  # Stage time offsets, length = stages
    Gamma: list[float]  # Gamma sums, length = stages
    NewF: list[bool]  # Whether stage i needs a fresh F evaluation
    ELO: float  # Estimator of local order (main + embedded + 1)

format_eqn(eqn_str, species_list, state_var='state', use_parentheses=True, keep_env_symbols=False, temperature=300.0, air_density=2.4476e+19)

Convert a symbolic ODE expression string into a C++ code string.

Parameters

eqn_str : str SymPy-parseable expression string (e.g., from KPP rate law). species_list : list Ordered species definitions used to map C_X symbols to state indices. state_var : str Name of the state array variable in generated C++ code. use_parentheses : bool If True, emit state(idx); if False, emit state_idx. keep_env_symbols : bool If False (default), environmental parameters (Temp, RH) are substituted with constants (Temp=300.0). This produces isothermal code suitable only for constant-temperature simulations. If True, Temp and RH are emitted as C++ variable references, and the generated function MUST accept temp/rh as parameters.

WARNING

The current default (keep_env_symbols=False) silently folds temperature to 300 K. For temperature-dependent chemistry, callers MUST pass keep_env_symbols=True and ensure the generated C++ function accepts temp/rh as parameters.

.. deprecated:: planned for MKPP 2.0 The default will be inverted to keep_env_symbols=True for physical correctness. This requires all generated function signatures to accept a temp parameter unconditionally.

Source code in src/mkpp/format_eqn.py 
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
def format_eqn(
    eqn_str,
    species_list,
    state_var="state",
    use_parentheses=True,
    keep_env_symbols=False,
    temperature: float = 300.0,
    air_density: float = 2.4476e19,
):
    """Convert a symbolic ODE expression string into a C++ code string.

    Parameters
    ----------
    eqn_str : str
        SymPy-parseable expression string (e.g., from KPP rate law).
    species_list : list
        Ordered species definitions used to map C_X symbols to state indices.
    state_var : str
        Name of the state array variable in generated C++ code.
    use_parentheses : bool
        If True, emit ``state(idx)``; if False, emit ``state_idx``.
    keep_env_symbols : bool
        If False (default), environmental parameters (Temp, RH) are substituted
        with constants (Temp=300.0). This produces isothermal code suitable only
        for constant-temperature simulations.
        If True, Temp and RH are emitted as C++ variable references, and the
        generated function MUST accept temp/rh as parameters.

    WARNING
    -------
    The current default (keep_env_symbols=False) silently folds temperature to
    300 K. For temperature-dependent chemistry, callers MUST pass
    keep_env_symbols=True and ensure the generated C++ function accepts temp/rh
    as parameters.

    .. deprecated:: planned for MKPP 2.0
        The default will be inverted to keep_env_symbols=True for physical
        correctness. This requires all generated function signatures to accept
        a ``temp`` parameter unconditionally.
    """
    import sympy as sp

    # 1. Clean up double negatives and malformed trailing dots on floats
    s = str(eqn_str).replace("--", "+").replace("^+", "^").replace("**+", "**")
    s = re.sub(r"(\d+\.\d+)\.+", r"\1", s)
    if s == "0":
        return "0.0"

    # 2. Try to use SymPy's C-code generator for robust math formatting
    try:
        expr = sp.sympify(s)
        # Substitute legacy KPP dummy vars to 1.0 before C-code generation
        subs_dict = {
            # Legacy KPP dummy/fixed species (not real state variables)
            sp.Symbol("C_DummyCH4"): 1.0,
            sp.Symbol("C_DummyNMVOC"): 1.0,
            sp.Symbol("C_FixedOH"): 1.0,
            sp.Symbol("C_FixedCl"): 1.0,
            sp.Symbol("S_a"): 1.0,
            sp.Symbol("v_gas"): 1.0,
            # M_density is used only when a mechanism has no explicit AIR or
            # M species. It must come from the supplied compilation
            # environment, never from a unit-specific legacy constant.
            sp.Symbol("M_density"): air_density,
        }

        # When keep_env_symbols is False (default), substitute environmental
        # parameters with constants for backward compatibility. When True,
        # leave Temp and RH as C variable references (used when equilibrium
        # reactions provide temp/rh as explicit function parameters).
        if not keep_env_symbols:
            # Environmental parameters (not species concentrations)
            # NOTE: SUN is NOT substituted — photolysis rates are runtime J-values from Cloud-J
            subs_dict[sp.Symbol("TEMP")] = temperature
            subs_dict[sp.Symbol("temp")] = temperature
            subs_dict[sp.Symbol("Temp")] = temperature

        expr = expr.subs(subs_dict)
        s = sp.ccode(expr)
        s = _fold_numeric_falloff_powers(s)
    except Exception:
        # Fallback to regex if sympy fails
        s = re.sub(r"([a-zA-Z0-9_\(\)\.\+\-\*\/]+)\*\*(\-?\d+\.\d+|\-?\d+)", r"pow(\1, \2)", s)
        if not keep_env_symbols:
            s = s.replace("Temp", str(temperature))
        s = s.replace("S_a", "1.0")
        s = s.replace("v_gas", "1.0")
        s = _fold_numeric_falloff_powers(s)

    # 3. Map the C_X species symbols from the SymPy AST directly into the state indices or variables.
    sorted_sp = sorted(list(enumerate(species_list)), key=lambda x: len(x[1].name), reverse=True)
    for idx_s, spec in sorted_sp:
        if use_parentheses:
            repl = f"{state_var}({idx_s})"
        else:
            repl = f"{state_var}_{idx_s}"
        s = re.sub(r"\bC_" + spec.name + r"(?!\w)", repl, s)

    # 4. Map J_<idx> photolysis symbols to the jvals array (Cloud-J runtime input)
    s = re.sub(r"\bJ_(\d+)\b", r"jvals[\1]", s)

    # 5. Map Rate_<idx> symbols (PHASE_CHANGE, TUNNELING) to jvals array
    # These are externally-provided rates from host model thermodynamic solvers (e.g., ISORROPIA)
    s = re.sub(r"\bRate_(\d+)\b", r"jvals[\1]", s)

    s = _strength_reduce_squares(s)
    s = _clean_float_literals(s)

    return s

generate_headers(mech, out_dir='src/solvers', suffix='', solver_name='ros3', adjoint=False, generate_host_api=False, simd_backend='native', emit_reference_backend=False)

Emit the Kokkos headers and manifest artifact.

Source code in src/mkpp/codegen.py 
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
def generate_headers(
    mech: MechanismDefinition,
    out_dir: str = "src/solvers",
    suffix: str = "",
    solver_name: str = "ros3",
    adjoint: bool = False,
    generate_host_api: bool = False,
    simd_backend: str = "native",
    emit_reference_backend: bool = False,
) -> dict[str, str]:
    """Emit the Kokkos headers and manifest artifact."""
    if not mech or not mech.species:
        raise ValueError("Cannot generate headers for empty mechanism")

    out_path = Path(out_dir)
    out_path.mkdir(parents=True, exist_ok=True)

    # Sensitivity code is always emitted into the generated header, guarded by
    # MKPP_ENABLE_ADJOINT.  That keeps a mechanism's chemistry artifact stable
    # while allowing CMake to preprocess adjoint/TLM code away in forward-only
    # builds.  ``adjoint`` remains accepted for CLI/API compatibility.
    context = build_template_context(mech, solver_name, adjoint=True, simd_backend=simd_backend)
    # Add suffix to context for filename generation
    context["suffix"] = suffix

    # 2. Render the header via Jinja2 template engine
    engine = TemplateEngine()
    header_text = engine.render("header.j2", context)

    # 3. Write rendered output to the same file path as before
    header_path = out_path / f"{mech.name}{suffix}.hpp"
    with open(header_path, "w") as f:
        f.write(header_text)

    compiled_sources = []
    # Keep the public include at the output root while grouping its compiled
    # implementation units under a mechanism-specific directory.
    compiled_path = out_path / f"{mech.name}{suffix}"
    compiled_path.mkdir(parents=True, exist_ok=True)
    for kernel, chunks in (("rates", context["compiled_rate_chunks"]), ("jacobian", context["compiled_jacobian_chunks"])):
        source_path = compiled_path / f"{kernel}.cpp"
        rendered_chunks = []
        for index, expressions in enumerate(chunks):
            source_context = dict(context)
            source_context.update({"compiled_kernel": kernel, "compiled_chunk_index": index, "compiled_expressions": expressions})
            rendered_chunks.append(engine.render("compiled_kernel_chunk.cpp.j2", source_context))
        with open(source_path, "w") as f:
            f.write("\n".join(rendered_chunks))
        for stale_path in compiled_path.glob(f"{kernel}_*.cpp"):
            stale_path.unlink()
        compiled_sources.append(str(source_path))

    if emit_reference_backend:
        for index, expressions in enumerate(context["compiled_lu_chunks"]):
            source_context = dict(context)
            source_context.update({"compiled_chunk_index": index, "compiled_expressions": expressions})
            source_path = compiled_path / f"factorize_{index}.cpp"
            with open(source_path, "w") as f:
                f.write(engine.render("compiled_factorize_chunk.cpp.j2", source_context))
            compiled_sources.append(str(source_path))
    else:
        for source_path in compiled_path.glob("factorize_*.cpp"):
            source_path.unlink()

    for kernel in ("solve",):
        source_path = compiled_path / f"{kernel}.cpp"
        with open(source_path, "w") as f:
            f.write(engine.render(f"compiled_{kernel}.cpp.j2", context))
        compiled_sources.append(str(source_path))

    for kernel in ("supernodal_factorize", "supernodal_solve"):
        source_path = compiled_path / f"{kernel}.cpp"
        with open(source_path, "w") as f:
            f.write(engine.render(f"{kernel}.cpp.j2", context))
        compiled_sources.append(str(source_path))

    results = {"header": str(header_path), "compiled_sources": compiled_sources}

    if generate_host_api:
        api_results = generate_host_api_headers(mech, out_dir=out_dir, solver_name=solver_name)
        results.update(api_results)

    # 4. Manifest metadata emission (unchanged logic)
    manifest = {
        "mechanism": mech.name,
        "aerosol_representation": mech.aerosol_representation.value,
        "simd_backend": simd_backend,
        "checksum": hashlib.sha256(mech.name.encode()).hexdigest(),
        "artifacts": [
            {"kind": "header", "file": header_path.name},
            *({"kind": "compiled_kernel_source", "file": str(Path(source).relative_to(out_path))} for source in compiled_sources),
            {"kind": "adjoint_tlm_record", "differentiable": True},
        ],
    }

    if getattr(mech, "host_interface", None) and mech.host_interface.arrays:
        manifest["host_interface"] = {
            arr.name: {
                "rank": arr.rank,
                "layout": arr.layout,
                "lifetime": "unmanaged_borrowed_from_host" if arr.ownership == "host" else "device_owned",
            }
            for arr in mech.host_interface.arrays
        }

    partition_meta = getattr(mech, "partition_metadata", None)
    if partition_meta:
        manifest["solver_partition"] = partition_meta

    manifest_path = out_path / f"{mech.name}_manifest.json"
    plan_json_path = compiled_path / "factorization_plan.json"
    plan_markdown_path = compiled_path / "factorization_plan.md"
    plan_json_path.write_text(engine.render("factorization_plan.json.j2", context), encoding="utf-8")
    plan_markdown_path.write_text(engine.render("factorization_plan.md.j2", context), encoding="utf-8")
    manifest["factorization_plan"] = {
        "json": str(plan_json_path.relative_to(out_path)),
        "markdown": str(plan_markdown_path.relative_to(out_path)),
    }
    with open(manifest_path, "w") as f:
        json.dump(manifest, f, indent=2)

    results["manifest"] = str(manifest_path)
    results["factorization_plan"] = str(plan_json_path)
    results["factorization_plan_markdown"] = str(plan_markdown_path)
    return results

generate_host_api_headers(mech, out_dir='src/solvers', solver_name='ros3')

Emit the C, C++, and Fortran host API headers and wrappers for a mechanism.

Source code in src/mkpp/codegen.py 
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
def generate_host_api_headers(
    mech: MechanismDefinition,
    out_dir: str = "src/solvers",
    solver_name: str = "ros3",
) -> dict[str, str]:
    """Emit the C, C++, and Fortran host API headers and wrappers for a mechanism."""
    if not mech or not mech.species:
        raise ValueError("Cannot generate host API headers for empty mechanism")

    out_path = Path(out_dir)
    out_path.mkdir(parents=True, exist_ok=True)

    context = build_template_context(mech, solver_name=solver_name)
    engine = TemplateEngine()

    rendered = {
        "c_header": ("mkpp.h", engine.render("host_api/mkpp.h.j2", context)),
        "fortran_module": ("mkpp_mod.f90", engine.render("host_api/mkpp_mod.f90.j2", context)),
        "cpp_header": ("mkpp.hpp", engine.render("host_api/mkpp.hpp.j2", context)),
        "c_api_source": ("mkpp_c_api.cpp", engine.render("host_api/mkpp_c_api.cpp.j2", context)),
    }

    results = {}
    for key, (filename, content) in rendered.items():
        file_path = out_path / filename
        with open(file_path, "w") as f:
            f.write(content)
        results[key] = str(file_path)

    return results

get_A(tableau, i, j)

Get A(i,j) from row-wise lower-triangular storage. i,j are 1-indexed; i > j.

Source code in src/mkpp/rosenbrock.py 
137
138
139
def get_A(tableau: RosenbrockTableau, i: int, j: int) -> float:
    """Get A(i,j) from row-wise lower-triangular storage. i,j are 1-indexed; i > j."""
    return tableau.A[(i - 1) * (i - 2) // 2 + j - 1]

get_C(tableau, i, j)

Get C(i,j) from row-wise lower-triangular storage. i,j are 1-indexed; i > j.

Source code in src/mkpp/rosenbrock.py 
142
143
144
def get_C(tableau: RosenbrockTableau, i: int, j: int) -> float:
    """Get C(i,j) from row-wise lower-triangular storage. i,j are 1-indexed; i > j."""
    return tableau.C[(i - 1) * (i - 2) // 2 + j - 1]