diff --git a/.github/workflows/Benchmarks.yml b/.github/workflows/Benchmarks.yml index 361b98844..e3bb1f698 100644 --- a/.github/workflows/Benchmarks.yml +++ b/.github/workflows/Benchmarks.yml @@ -24,3 +24,6 @@ jobs: with: julia-version: ${{ matrix.version }} bench-on: ${{ github.event.pull_request.head.sha }} + extra-pkgs: | + https://github.com/PalmStudio/XPalm.jl#main + https://github.com/VEZY/PlantBiophysics.jl#master diff --git a/benchmark/test-xpalm.jl b/benchmark/test-xpalm.jl index 13d891a2b..e3d9f7c98 100644 --- a/benchmark/test-xpalm.jl +++ b/benchmark/test-xpalm.jl @@ -47,6 +47,8 @@ function xpalm_default_param_convert_outputs(sim_outputs) end +println(Pkg.status("XPalm")) + #=@testset "XPalm simple test" begin # default number of seconds is 5 b_XP = @benchmark xpalm_default_param_run() seconds = 120 diff --git a/src/PlantSimEngine.jl b/src/PlantSimEngine.jl index 1151b8f7a..b2f385d71 100644 --- a/src/PlantSimEngine.jl +++ b/src/PlantSimEngine.jl @@ -106,6 +106,7 @@ include("time/runtime/meteo_sampling.jl") # Simulation: include("run.jl") +include("compiled_model.jl") # Fitting include("evaluation/fit.jl") @@ -137,6 +138,7 @@ export inputs, outputs, variables, convert_outputs export timespec, output_policy, timestep_hint, meteo_hint export input_bindings, meteo_bindings, meteo_window, output_routing, model_scope export run! +export compile_model, write_compiled_model export fit # Re-exporting PlantMeteo main functions: diff --git a/src/compiled_model.jl b/src/compiled_model.jl new file mode 100644 index 000000000..5712f1020 --- /dev/null +++ b/src/compiled_model.jl @@ -0,0 +1,986 @@ +""" + compile_model(mapping; function_name=:compiled_model!, mode=:readable) + compile_model(model_list; function_name=:compiled_model!, mode=:readable) + +Generate a Julia script containing one merged model function for a single-scale +model mapping. The generated function keeps `models`, `status`, `meteo`, +`constants`, and `extra` as arguments, so model parameters still come from the +original mapping/model list. + +Hard dependency calls to `run!(models.process, ...)` are recursively replaced by +the source body of the called model. Soft dependencies are emitted in dependency +graph order. +""" +function compile_model(mapping::ModelMapping{SingleScale}; function_name::Symbol=:compiled_model!, mode::Symbol=:readable) + return compile_model(_modellist_from_model_mapping(mapping); function_name=function_name, mode=mode) +end + +function compile_model(::ModelMapping{MultiScale}; function_name::Symbol=:compiled_model!, mode::Symbol=:readable) + _validate_compilation_mode(mode) + error("`compile_model` currently supports single-scale `ModelMapping`s. Build a `GraphSimulation` first if you need MTG-specific execution context.") +end + +function compile_model(model_list::ModelList; function_name::Symbol=:compiled_model!, mode::Symbol=:readable) + _validate_compilation_mode(mode) + dep_graph = dep(model_list) + length(dep_graph.not_found) > 0 && error("Cannot compile model with missing dependencies: $(dep_graph.not_found)") + + nodes = _topological_soft_nodes(dep_graph) + ctx = _ModelCompilationContext(model_list.models, Dict{Tuple{Symbol,Symbol,DataType},_CompiledModelSource}(), Set{Module}(), Dict(:Default => 1), mode) + statements = Any[] + for node in nodes + call = _CompiledRunCall( + :(models.$(node.process)), + :models, + :status, + :meteo, + :constants, + :extra + ) + append!(statements, _inline_model_node(ctx, node, call, Set{Tuple{Symbol,Symbol}}())) + end + + body = Expr(:block, statements..., :(return nothing)) + fexpr = Expr( + :function, + Expr(:call, function_name, :models, :status, :meteo, :constants, :extra), + body + ) + + return string( + "# This file was generated by PlantSimEngine.compile_model.\n", + "# It expects `models` to be the model mapping/list used for compilation.\n\n", + "# Mode: $(ctx.mode) single-scale model.\n\n", + _compiled_module_imports(ctx.source_modules), + _sprint_expr(fexpr), + "\n" + ) +end + +function compile_model(sim::GraphSimulation; function_name::Symbol=:compiled_model!, mode::Symbol=:readable) + _validate_compilation_mode(mode) + dep_graph = dep(sim) + length(dep_graph.not_found) > 0 && error("Cannot compile model with missing dependencies: $(dep_graph.not_found)") + + nodes = _topological_soft_nodes(dep_graph) + ctx = _ModelCompilationContext(get_models(sim), Dict{Tuple{Symbol,Symbol,DataType},_CompiledModelSource}(), Set{Module}(), _initial_status_counts(sim), mode) + return is_multirate(sim) ? + _compile_multirate_graph_simulation(sim, nodes, ctx, function_name) : + _compile_same_rate_graph_simulation(sim, nodes, ctx, function_name) +end + +function _compile_same_rate_graph_simulation(sim::GraphSimulation, nodes::Vector{SoftDependencyNode}, ctx, function_name::Symbol) + signature = _compiled_model_signature_expr(_compiled_model_signature(sim)) + step_statements = Any[] + for node in nodes + append!(step_statements, _multiscale_node_loop(ctx, sim, node)) + end + + one_step_body = Expr(:block, step_statements...) + body = Expr( + :block, + :(statuses = PlantSimEngine.status(sim)), + :(models_by_scale = PlantSimEngine.get_models(sim)), + :(PlantSimEngine._assert_compiled_sim_compatible(sim, $signature, false)), + :(nsteps = PlantSimEngine.get_nsteps(meteo)), + Expr( + :if, + :(nsteps == 1), + Expr( + :block, + :(meteo_i = meteo), + one_step_body, + :(PlantSimEngine.save_results!(sim, 1)) + ), + Expr( + :block, + :(i = 1), + Expr( + :for, + :(meteo_i = PlantSimEngine.Tables.rows(meteo)), + Expr( + :block, + one_step_body, + :(PlantSimEngine.save_results!(sim, i)), + :(i += 1) + ) + ) + ) + ), + Expr( + :for, + :((organ, index) = sim.outputs_index), + :(resize!(PlantSimEngine.outputs(sim)[organ], index - 1)) + ), + :(return PlantSimEngine.outputs(sim)) + ) + fexpr = Expr(:function, Expr(:call, function_name, :sim, :meteo, :constants), body) + + return string( + "# This file was generated by PlantSimEngine.compile_model.\n", + "# It expects `sim` to be the initialized GraphSimulation used for compilation.\n", + "# Mode: same-rate multiscale MTG; compiler_mode: $(ctx.mode); variables are scale-scoped through each Status.\n", + "# Cross-scale sharing comes from initialized Refs and RefVectors in sim.statuses.\n", + "# Initial statuses by scale: $(_initial_status_counts(sim)).\n", + "# Processes by scale: $(_processes_by_scale(sim)).\n\n", + "# Model compatibility signature: $(_compiled_model_signature(sim)).\n\n", + _compiled_module_imports(ctx.source_modules), + _sprint_expr(fexpr), + "\n" + ) +end + +function _compile_multirate_graph_simulation(sim::GraphSimulation, nodes::Vector{SoftDependencyNode}, ctx, function_name::Symbol) + signature = _compiled_model_signature_expr(_compiled_model_signature(sim)) + setup_statements = Any[] + step_statements = Any[] + for node in nodes + locals = _multirate_node_locals(node) + append!(setup_statements, _multirate_node_setup(node, locals)) + append!(step_statements, _multirate_node_loop(ctx, sim, node, locals)) + end + + one_step_body = Expr( + :block, + :(t = PlantSimEngine._time_from_step(i, timeline)), + step_statements..., + :(PlantSimEngine.update_requested_outputs!(sim, t)), + :(PlantSimEngine.save_results!(sim, i)) + ) + body = Expr( + :block, + :(PlantSimEngine._validate_meteo_duration(meteo)), + :(dep_graph = PlantSimEngine.dep(sim)), + :(statuses = PlantSimEngine.status(sim)), + :(models_by_scale = PlantSimEngine.get_models(sim)), + :(PlantSimEngine._assert_compiled_sim_compatible(sim, $signature, true)), + :(compiled_nodes = PlantSimEngine._compiled_dependency_nodes_by_key(sim)), + :(timeline = PlantSimEngine._timeline_context(meteo)), + :(meteo_sampler = PlantSimEngine._prepare_meteo_sampler(meteo)), + :(runtime_clock_rows = PlantSimEngine._runtime_clock_rows(sim, timeline, dep_graph)), + :(PlantSimEngine._validate_meteo_derived_timestep_requirements!(runtime_clock_rows, timeline)), + :(PlantSimEngine._warn_if_no_model_runs_at_base_timestep(runtime_clock_rows, timeline)), + :(PlantSimEngine.validate_canonical_publishers(sim)), + :(PlantSimEngine.prepare_output_requests!(sim, tracked_outputs, timeline)), + :(PlantSimEngine.configure_temporal_buffers!(sim, timeline)), + setup_statements..., + :(nsteps = PlantSimEngine.get_nsteps(meteo)), + Expr( + :if, + :(nsteps == 1), + Expr( + :block, + :(i = 1), + :(meteo_i = meteo), + one_step_body + ), + Expr( + :block, + :(i = 1), + Expr( + :for, + :(meteo_i = PlantSimEngine.Tables.rows(meteo)), + Expr( + :block, + one_step_body, + :(i += 1) + ) + ) + ) + ), + Expr( + :for, + :((organ, index) = sim.outputs_index), + :(resize!(PlantSimEngine.outputs(sim)[organ], index - 1)) + ), + Expr( + :if, + :return_requested_outputs, + :(return PlantSimEngine.outputs(sim), PlantSimEngine.collect_outputs(sim; sink=requested_outputs_sink)), + :(return PlantSimEngine.outputs(sim)) + ) + ) + fexpr = Expr( + :function, + Expr( + :call, + function_name, + Expr( + :parameters, + Expr(:kw, :tracked_outputs, :nothing), + Expr(:kw, :return_requested_outputs, false), + Expr(:kw, :requested_outputs_sink, :(PlantSimEngine.DataFrames.DataFrame)) + ), + :sim, + :meteo, + :constants + ), + body + ) + + return string( + "# This file was generated by PlantSimEngine.compile_model.\n", + "# It expects `sim` to be the initialized GraphSimulation used for compilation.\n", + "# Mode: multi-rate multiscale MTG; compiler_mode: $(ctx.mode); variables are scale-scoped through each Status.\n", + "# Temporal inputs, meteo sampling, output routing, scopes, and online exports use PlantSimEngine runtime helpers.\n", + "# Cross-scale sharing comes from initialized Refs and RefVectors in sim.statuses.\n", + "# Initial statuses by scale: $(_initial_status_counts(sim)).\n", + "# Processes by scale: $(_processes_by_scale(sim)).\n\n", + "# Model compatibility signature: $(_compiled_model_signature(sim)).\n\n", + _compiled_module_imports(ctx.source_modules), + _sprint_expr(fexpr), + "\n" + ) +end + +""" + write_compiled_model(path, mapping; function_name=:compiled_model!, mode=:readable) + +Write the script generated by [`compile_model`](@ref) to `path` and return the +path. +""" +function write_compiled_model(path::AbstractString, mapping; function_name::Symbol=:compiled_model!, mode::Symbol=:readable) + open(path, "w") do io + write(io, compile_model(mapping; function_name=function_name, mode=mode)) + end + return path +end + +struct _CompiledModelSource + method::Method + expression::Expr + body::Expr + parameters::Vector{Pair{Union{Nothing,Symbol},Any}} +end + +struct _CompiledRunCall + model + models + status + meteo + constants + extra +end + +struct _ModelCompilationContext{M} + models::M + source_cache::Dict{Tuple{Symbol,Symbol,DataType},_CompiledModelSource} + source_modules::Set{Module} + initial_status_counts::Dict{Symbol,Int} + mode::Symbol +end + +function _validate_compilation_mode(mode::Symbol) + mode in (:readable, :fast) || error("Unsupported compiled model mode `$mode`. Expected `:readable` or `:fast`.") + return mode +end + +function _topological_soft_nodes(dep_graph::DependencyGraph) + traversal_order = traverse_dependency_graph(dep_graph, false) + pending = Set(traversal_order) + ordered = SoftDependencyNode[] + + while !isempty(pending) + progressed = false + for node in traversal_order + node in pending || continue + parents = isnothing(node.parent) ? SoftDependencyNode[] : node.parent + if all(parent -> !(parent in pending), parents) + push!(ordered, node) + delete!(pending, node) + progressed = true + end + end + progressed || error("Cannot compile cyclic dependency graph.") + end + + return ordered +end + +function _inline_model_node(ctx::_ModelCompilationContext, node::AbstractDependencyNode, call::_CompiledRunCall, stack::Set{Tuple{Symbol,Symbol}}) + node_key = (node.scale, node.process) + node_key in stack && error("Cyclic hard dependency while compiling process `$(node.process)` at scale `$(node.scale)`.") + + source_key = (node.scale, node.process, typeof(node.value)) + source = get!(ctx.source_cache, source_key) do + _compiled_model_source(ctx, node.value) + end + compiled_body = try + body_call = _body_call(ctx, node, call) + bindings = _argument_bindings(source, body_call) + rewritten = _rewrite_run_body(ctx, source.body, bindings, push!(copy(stack), node_key), node.scale) + rewritten = _rewrite_current_model_reference(rewritten, body_call.models, node.process, body_call.model) + _ensure_no_uninlined_run_calls(rewritten, node, source) + _strip_return_expr(rewritten) + catch err + err isa ErrorException || rethrow() + error( + "Failed to compile process `$(node.process)` at scale `$(node.scale)` from `$(source.method)` ", + "defined in $(String(source.method.file)):$(source.method.line). ", + err.msg + ) + end + + statements = Any[ + LineNumberNode(0, Symbol(_model_metadata_comment(ctx, node, source))), + ] + if ctx.mode == :readable + push!(statements, LineNumberNode(0, Symbol(_model_io_comment(node)))) + append!(statements, _readable_block_bindings(node, call)) + else + push!(statements, :(model = $(call.model))) + end + append!(statements, _block_statements(compiled_body)) + + return Any[Expr(:let, Expr(:block), Expr(:block, statements...))] +end + +function _body_call(ctx::_ModelCompilationContext, node::AbstractDependencyNode, call::_CompiledRunCall) + ctx.mode == :readable || return call + prefix = _readable_node_prefix(node) + return _CompiledRunCall( + Symbol(prefix, "_model"), + Symbol(prefix, "_models"), + Symbol(prefix, "_status"), + Symbol(prefix, "_meteo"), + :constants, + call.extra, + ) +end + +function _readable_block_bindings(node::AbstractDependencyNode, call::_CompiledRunCall) + prefix = _readable_node_prefix(node) + local_model = Symbol(prefix, "_model") + local_models = Symbol(prefix, "_models") + local_status = Symbol(prefix, "_status") + local_meteo = Symbol(prefix, "_meteo") + return Any[ + Expr(:local, local_model, local_models, local_status, local_meteo), + :($local_model = $(call.model)), + :($local_models = $(call.models)), + :($local_status = $(call.status)), + :($local_meteo = $(call.meteo)), + ] +end + +function _readable_node_prefix(node::AbstractDependencyNode) + scale = node.scale == :Default ? "" : string(_identifier_part(node.scale), "_") + return Symbol(scale, _identifier_part(node.process)) +end + +function _rewrite_current_model_reference(expr, models_expr, process::Symbol, model_expr) + _is_current_model_reference(expr, models_expr, process) && return model_expr + expr isa Expr || return expr + return Expr(expr.head, (_rewrite_current_model_reference(arg, models_expr, process, model_expr) for arg in expr.args)...) +end + +function _is_current_model_reference(expr, models_expr, process::Symbol) + if Meta.isexpr(expr, :.) && length(expr.args) == 2 && expr.args[2] == QuoteNode(process) + return expr.args[1] == models_expr + elseif Meta.isexpr(expr, :call) && length(expr.args) == 3 && expr.args[1] in (:getfield, :getproperty) + return expr.args[2] == models_expr && _quote_value(expr.args[3]) == process + end + return false +end + +function _multiscale_node_loop(ctx::_ModelCompilationContext, sim::GraphSimulation, node::SoftDependencyNode) + call = _CompiledRunCall( + :((models_by_scale[$(QuoteNode(node.scale))]).$(node.process)), + :(models_by_scale[$(QuoteNode(node.scale))]), + :status, + :meteo_i, + :constants, + :sim + ) + body = _inline_model_node(ctx, node, call, Set{Tuple{Symbol,Symbol}}()) + statements = Any[ + LineNumberNode(0, Symbol("scale loop: $(node.scale) | process: $(node.process) | initial_statuses: $(length(status(sim)[node.scale]))")), + :(idx = 1), + :(idx_limit = length(statuses[$(QuoteNode(node.scale))])), + Expr( + :while, + :(idx <= idx_limit), + Expr( + :block, + :(status = statuses[$(QuoteNode(node.scale))][idx]), + body..., + :(idx += 1) + ) + ) + ] + return Any[Expr(:let, Expr(:block), Expr(:block, statements...))] +end + +function _multirate_node_locals(node::SoftDependencyNode) + prefix = string("_mr_", _identifier_part(node.scale), "_", _identifier_part(node.process)) + return ( + node=Symbol(prefix, "_node"), + model=Symbol(prefix, "_model"), + model_spec=Symbol(prefix, "_model_spec"), + model_clock=Symbol(prefix, "_model_clock"), + ) +end + +function _multirate_node_setup(node::SoftDependencyNode, locals) + model_expr = :((models_by_scale[$(QuoteNode(node.scale))]).$(node.process)) + return Any[ + LineNumberNode(0, Symbol("multirate setup: $(node.scale) | process: $(node.process)")), + :($(locals.node) = compiled_nodes[($(QuoteNode(node.scale)), $(QuoteNode(node.process)))]), + :($(locals.model) = $model_expr), + :($(locals.model_spec) = PlantSimEngine._model_spec_for_process(sim, $(QuoteNode(node.scale)), $(QuoteNode(node.process)))), + :($(locals.model_clock) = PlantSimEngine._model_clock($(locals.model_spec), $(locals.model), timeline)), + ] +end + +function _multirate_node_loop(ctx::_ModelCompilationContext, sim::GraphSimulation, node::SoftDependencyNode, locals) + model_expr = locals.model + call = _CompiledRunCall( + model_expr, + :(models_by_scale[$(QuoteNode(node.scale))]), + :status, + :meteo_for_model, + :constants, + :sim + ) + body = _inline_model_node(ctx, node, call, Set{Tuple{Symbol,Symbol}}()) + statements = Any[ + LineNumberNode(0, Symbol("multirate scale loop: $(node.scale) | process: $(node.process) | initial_statuses: $(length(status(sim)[node.scale]))")), + Expr( + :if, + :(PlantSimEngine._should_run_at_time($(locals.model_clock), t)), + Expr( + :block, + :(idx = 1), + :(idx_limit = length(statuses[$(QuoteNode(node.scale))])), + Expr( + :while, + :(idx <= idx_limit), + Expr( + :block, + :(status = statuses[$(QuoteNode(node.scale))][idx]), + :(PlantSimEngine.resolve_inputs_from_temporal_state!(sim, $(locals.node), status, t, $(locals.model_spec), timeline)), + :(meteo_for_model = PlantSimEngine._sample_meteo_for_model(meteo_sampler, meteo_i, i, $(locals.model_clock), $(locals.model_spec))), + body..., + :(PlantSimEngine.update_temporal_state_outputs!(sim, $(locals.node), $(locals.model_spec), status, t)), + :(idx += 1) + ) + ) + ), + nothing + ) + ] + return Any[Expr(:let, Expr(:block), Expr(:block, statements...))] +end + +function _identifier_part(x) + raw = string(x) + chars = map(collect(raw)) do c + isletter(c) || isdigit(c) ? c : '_' + end + cleaned = String(chars) + isempty(cleaned) && return "unnamed" + return isdigit(first(cleaned)) ? string("_", cleaned) : cleaned +end + +function _compiled_dependency_nodes_by_key(sim::GraphSimulation) + nodes = Dict{Tuple{Symbol,Symbol},SoftDependencyNode}() + for node in _topological_soft_nodes(dep(sim)) + nodes[(node.scale, node.process)] = node + end + return nodes +end + +function _compiled_model_signature(sim::GraphSimulation) + entries = Tuple{Symbol,Symbol,String}[] + for scale in sort!(collect(keys(get_models(sim))); by=string) + models_at_scale = get_models(sim)[scale] + for process in sort!(collect(keys(models_at_scale)); by=string) + push!(entries, (scale, process, string(typeof(getproperty(models_at_scale, process))))) + end + end + return Tuple(entries) +end + +function _compiled_model_signature_expr(signature) + return Expr(:tuple, (Expr(:tuple, QuoteNode(scale), QuoteNode(process), type_name) for (scale, process, type_name) in signature)...) +end + +function _assert_compiled_sim_compatible(sim::GraphSimulation, expected_signature, expected_multirate::Bool) + is_multirate(sim) == expected_multirate || error( + "Compiled model was generated for ", + expected_multirate ? "a multi-rate" : "a same-rate", + " GraphSimulation, but received ", + is_multirate(sim) ? "a multi-rate" : "a same-rate", + " GraphSimulation." + ) + actual_signature = _compiled_model_signature(sim) + actual_signature == expected_signature && return nothing + + error( + "Compiled model was generated for a different model mapping. ", + "Expected signature: $(expected_signature). ", + "Received signature: $(actual_signature)." + ) +end + +function _compiled_model_source(ctx::_ModelCompilationContext, model) + method = _run_method_for_model(model) + push!(ctx.source_modules, method.module) + fexpr = _method_expr_from_source(method) + return _CompiledModelSource(method, fexpr, _method_body_from_expr(method, fexpr), _method_parameters(fexpr)) +end + +function _argument_bindings(source::_CompiledModelSource, call::_CompiledRunCall) + values = Any[call.model, call.models, call.status, call.meteo, call.constants, call.extra] + bindings = Dict{Symbol,Any}() + for (i, parameter) in enumerate(source.parameters) + name, default = first(parameter), last(parameter) + isnothing(name) && continue + if i <= length(values) + bindings[name] = values[i] + elseif !isnothing(default) + bindings[name] = default + else + error("Cannot inline `$(source.method)`: missing argument `$name` and no default value was found.") + end + end + return bindings +end + +function _model_metadata_comment(ctx::_ModelCompilationContext, node::AbstractDependencyNode, source::_CompiledModelSource) + initial_statuses = get(ctx.initial_status_counts, node.scale, 0) + return join( + ( + "model: $(typeof(node.value))", + "process: $(node.process)", + "scale: $(node.scale)", + "initial_statuses: $initial_statuses", + "module: $(source.method.module)", + "source: $(String(source.method.file)):$(source.method.line)", + "method: $(source.method)" + ), + " | " + ) +end + +function _model_io_comment(node::AbstractDependencyNode) + return join( + ( + "inputs: $(_variable_list_comment(inputs_(node.value)))", + "outputs: $(_variable_list_comment(outputs_(node.value)))", + ), + " | " + ) +end + +function _variable_list_comment(vars) + names = collect(keys(vars)) + isempty(names) && return "none" + return join(string.(names), ", ") +end + +function _compiled_module_imports(modules::Set{Module}) + imports = String[ + "using PlantSimEngine\n", + "using MultiScaleTreeGraph\n", + "using PlantMeteo\n", + "using Statistics\n", + ] + for mod in sort!(collect(modules), by=string) + mod in (PlantSimEngine, Main) && continue + push!(imports, "using $(_module_path(mod))\n") + end + push!(imports, "\n") + return join(imports) +end + +function _module_path(mod::Module) + names = Symbol[] + current = mod + while current !== Main + pushfirst!(names, nameof(current)) + parent = parentmodule(current) + parent === current && break + current = parent + end + return join(string.(names), ".") +end + +function _initial_status_counts(sim::GraphSimulation) + return Dict(scale => length(statuses) for (scale, statuses) in status(sim)) +end + +function _processes_by_scale(sim::GraphSimulation) + return Dict(scale => collect(keys(models)) for (scale, models) in get_models(sim)) +end + +function _run_method_for_model(model) + model_type = typeof(model) + candidates = Method[] + for method in methods(run!) + sig = Base.unwrap_unionall(method.sig) + sig <: Tuple || continue + length(sig.parameters) >= 2 || continue + model_arg = sig.parameters[2] + model_arg isa Type || continue + model_type <: model_arg || continue + push!(candidates, method) + end + + isempty(candidates) && error("No `run!` method found for model type `$(model_type)`.") + exact = filter(candidates) do method + sig = Base.unwrap_unionall(method.sig) + sig.parameters[2] == model_type + end + if !isempty(exact) + max_arity = maximum(method -> length(Base.unwrap_unionall(method.sig).parameters), exact) + exact_max = filter(method -> length(Base.unwrap_unionall(method.sig).parameters) == max_arity, exact) + length(exact_max) == 1 && return only(exact_max) + end + + non_any = filter(candidates) do method + sig = Base.unwrap_unionall(method.sig) + sig.parameters[2] != Any + end + if !isempty(non_any) + max_arity = maximum(method -> length(Base.unwrap_unionall(method.sig).parameters), non_any) + non_any_max = filter(method -> length(Base.unwrap_unionall(method.sig).parameters) == max_arity, non_any) + length(non_any_max) == 1 && return only(non_any_max) + end + + length(candidates) == 1 && return only(candidates) + + error( + "Several `run!` methods match model type `$(model_type)`. ", + "`compile_model` needs a single concrete model method." + ) +end + +function _method_body_from_expr(method::Method, fexpr::Expr) + body = if Meta.isexpr(fexpr, :function) + fexpr.args[2] + elseif Meta.isexpr(fexpr, :(=)) + fexpr.args[2] + else + error("Unsupported method expression for `$method`.") + end + return _clean_source_expr(body) +end + +function _method_expr_from_source(method::Method) + file = String(method.file) + line = method.line + isfile(file) || error("Cannot find source file for `$method`: $file") + lines = readlines(file) + line <= length(lines) || error("Invalid source location for `$method`: $file:$line") + + buffer = IOBuffer() + last_error = nothing + for i in line:length(lines) + println(buffer, lines[i]) + source = String(take!(copy(buffer))) + try + expr = Meta.parse(source) + if Meta.isexpr(expr, :function) || Meta.isexpr(expr, :(=)) + return expr + end + catch err + last_error = err + end + end + error("Failed to parse source for `$method` from $file:$line. Last parser error: $last_error") +end + +function _clean_source_expr(expr) + if expr isa LineNumberNode + return nothing + elseif Meta.isexpr(expr, :block) + return Expr(:block, filter(!isnothing, Any[_clean_source_expr(arg) for arg in expr.args])...) + elseif expr isa Expr + return Expr(expr.head, filter(!isnothing, Any[_clean_source_expr(arg) for arg in expr.args])...) + else + return expr + end +end + +function _function_signature(expr::Expr) + if Meta.isexpr(expr, :function) + return expr.args[1] + elseif Meta.isexpr(expr, :(=)) + return expr.args[1] + end + error("Unsupported function expression.") +end + +function _call_arguments(sig) + if Meta.isexpr(sig, :call) + return sig.args[2:end] + elseif Meta.isexpr(sig, :where) + return _call_arguments(sig.args[1]) + end + return Any[] +end + +function _method_parameters(fexpr::Expr) + args = _call_arguments(_function_signature(fexpr)) + return Pair{Union{Nothing,Symbol},Any}[_argument_name(arg) => _argument_default(arg) for arg in args] +end + +function _argument_default(arg) + if Meta.isexpr(arg, :(=)) || Meta.isexpr(arg, :kw) + return arg.args[2] + end + return nothing +end + +function _argument_name(arg) + if Meta.isexpr(arg, :(=)) + return _argument_name(arg.args[1]) + end + return _argument_name_without_default(arg) +end + +function _argument_name_without_default(arg) + arg isa Symbol && return arg + if Meta.isexpr(arg, :(::)) + return length(arg.args) == 2 && arg.args[1] isa Symbol ? arg.args[1] : nothing + elseif Meta.isexpr(arg, :kw) + return _argument_name(arg.args[1]) + end + return nothing +end + +function _rewrite_run_body(ctx::_ModelCompilationContext, expr, bindings::Dict{Symbol,Any}, stack::Set{Tuple{Symbol,Symbol}}, current_scale::Symbol) + if expr isa Symbol + return get(bindings, expr, expr) + elseif expr isa QuoteNode || expr isa LineNumberNode + return expr + elseif expr isa Expr + if Meta.isexpr(expr, :block) + local_bindings = copy(bindings) + return Expr(:block, (_rewrite_block_statement(ctx, arg, local_bindings, stack, current_scale) for arg in expr.args)...) + elseif Meta.isexpr(expr, :(=)) && expr.args[1] isa Symbol + rewritten_rhs = _rewrite_run_body(ctx, expr.args[2], bindings, stack, current_scale) + _update_binding!(bindings, expr.args[1], rewritten_rhs) + return Expr(:(=), expr.args[1], rewritten_rhs) + elseif Meta.isexpr(expr, :call) + rewritten_f = _rewrite_run_body(ctx, expr.args[1], bindings, stack, current_scale) + if _is_run_function(rewritten_f) + rewritten_call = Expr(:call, rewritten_f, expr.args[2:end]...) + return _inline_run_call(ctx, rewritten_call, bindings, stack, current_scale) + end + rewritten_args = Any[_rewrite_run_body(ctx, arg, bindings, stack, current_scale) for arg in expr.args[2:end]] + return Expr(:call, rewritten_f, rewritten_args...) + end + return Expr(expr.head, (_rewrite_run_body(ctx, arg, bindings, stack, current_scale) for arg in expr.args)...) + end + return expr +end + +function _rewrite_block_statement(ctx::_ModelCompilationContext, expr, bindings::Dict{Symbol,Any}, stack::Set{Tuple{Symbol,Symbol}}, current_scale::Symbol) + if expr isa Expr && Meta.isexpr(expr, :(=)) && expr.args[1] isa Symbol + rewritten_rhs = _rewrite_run_body(ctx, expr.args[2], bindings, stack, current_scale) + _update_binding!(bindings, expr.args[1], rewritten_rhs) + return Expr(:(=), expr.args[1], rewritten_rhs) + end + return _rewrite_run_body(ctx, expr, bindings, stack, current_scale) +end + +function _update_binding!(bindings::Dict{Symbol,Any}, name::Symbol, value) + if _is_bindable_alias(value) + bindings[name] = value + else + delete!(bindings, name) + end + return value +end + +function _is_bindable_alias(value) + value isa Symbol && return true + value isa QuoteNode && return true + value isa LineNumberNode && return false + value isa Expr || return true + if Meta.isexpr(value, :call) && value.args[1] in (:getfield, :getproperty) + return all(_is_bindable_alias(arg) for arg in value.args[2:end]) + elseif Meta.isexpr(value, :(.)) + return all(_is_bindable_alias(arg) for arg in value.args) + end + return false +end + +function _substitute_bound_symbols(expr, bindings::Dict{Symbol,Any}) + if expr isa Symbol + return get(bindings, expr, expr) + elseif expr isa QuoteNode || expr isa LineNumberNode + return expr + elseif expr isa Expr + return Expr(expr.head, (_substitute_bound_symbols(arg, bindings) for arg in expr.args)...) + end + return expr +end + +function _inline_run_call(ctx::_ModelCompilationContext, expr::Expr, bindings::Dict{Symbol,Any}, stack::Set{Tuple{Symbol,Symbol}}, current_scale::Symbol) + call_args = Any[_substitute_bound_symbols(arg, bindings) for arg in expr.args[2:end]] + length(call_args) >= 4 || error("Cannot inline `run!` call with fewer than four runtime arguments: `$expr`.") + scale, process = _resolve_model_reference(call_args[1], current_scale) + model = _model_from_context(ctx, scale, process) + hard_node = HardDependencyNode( + model, + process, + NamedTuple(), + Int[], + scale, + inputs_(model), + outputs_(model), + nothing, + HardDependencyNode[] + ) + full_args = _complete_run_call_args(call_args) + call = _CompiledRunCall(full_args[1], full_args[2], full_args[3], full_args[4], full_args[5], full_args[6]) + return Expr(:block, _inline_model_node(ctx, hard_node, call, stack)...) +end + +function _ensure_no_uninlined_run_calls(expr, node::AbstractDependencyNode, source::_CompiledModelSource) + calls = _find_run_calls(expr) + isempty(calls) && return nothing + rendered = join(("`$(_sprint_expr(call))`" for call in calls), ", ") + error( + "Unsupported `run!` call shape remained after source rewriting in `$(source.method)`. ", + "Process `$(node.process)` at scale `$(node.scale)` still contains: $(rendered)" + ) +end + +function _find_run_calls(expr) + calls = Expr[] + _collect_run_calls!(calls, expr) + return calls +end + +function _collect_run_calls!(calls::Vector{Expr}, expr) + expr isa Expr || return calls + if Meta.isexpr(expr, :call) && _is_run_function(expr.args[1]) + push!(calls, expr) + end + for arg in expr.args + _collect_run_calls!(calls, arg) + end + return calls +end + +function _complete_run_call_args(args::Vector{Any}) + values = Any[args...] + length(values) < 5 && push!(values, :(nothing)) + length(values) < 6 && push!(values, :(nothing)) + length(values) == 6 || error("Cannot inline `run!` call with unsupported arity $(length(args)).") + return values +end + +function _model_from_context(ctx::_ModelCompilationContext, scale::Symbol, process::Symbol) + if ctx.models isa AbstractDict + haskey(ctx.models, scale) || error("Cannot inline hard dependency `$(process)`: scale `$(scale)` is not in the model mapping.") + hasproperty(ctx.models[scale], process) || error("Cannot inline hard dependency `$(process)` at scale `$(scale)`: no model found.") + return getproperty(ctx.models[scale], process) + end + scale == :Default || error("Cannot inline hard dependency `$(process)` at scale `$(scale)` in a single-scale model mapping.") + hasproperty(ctx.models, process) || error("Cannot inline hard dependency `$(process)`: no model found.") + return getproperty(ctx.models, process) +end + +function _resolve_model_reference(expr, current_scale::Symbol) + if Meta.isexpr(expr, :.) && length(expr.args) == 2 && expr.args[2] isa QuoteNode + process = expr.args[2].value + scale = _resolve_models_container_scale(expr.args[1], current_scale) + return scale, process + elseif Meta.isexpr(expr, :call) && length(expr.args) == 3 && expr.args[1] in (:getfield, :getproperty) + scale = _resolve_models_container_scale(expr.args[2], current_scale) + process = _quote_value(expr.args[3]) + process isa Symbol || error("Cannot inline hard dependency call. Expected a literal process symbol in `$expr`.") + return scale, process + end + error("Cannot inline hard dependency call. Unsupported model reference `$expr`.") +end + +function _resolve_models_container_scale(expr, current_scale::Symbol) + expr == :models && return current_scale + _is_readable_models_local(expr) && return current_scale + if Meta.isexpr(expr, :ref) && length(expr.args) == 2 + base, scale_expr = expr.args + scale = _quote_value(scale_expr) + if base == :models_by_scale || _is_sim_models_expr(base) + scale isa Symbol || error("Cannot inline hard dependency call. Expected a literal scale symbol in `$expr`.") + return scale + end + end + error("Cannot inline hard dependency call. Unsupported models container `$expr`.") +end + +function _is_readable_models_local(expr) + expr isa Symbol || return false + return endswith(String(expr), "_models") +end + +function _quote_value(expr) + expr isa QuoteNode && return expr.value + expr isa Symbol && return expr + return nothing +end + +function _is_sim_models_expr(expr) + return Meta.isexpr(expr, :.) && + length(expr.args) == 2 && + expr.args[1] == :sim && + expr.args[2] == QuoteNode(:models) +end + +_is_run_function(f) = f == :run! || f == GlobalRef(@__MODULE__, :run!) +function _is_run_function(f::Expr) + return Meta.isexpr(f, :.) && + length(f.args) == 2 && + f.args[1] in (:PlantSimEngine, Symbol(@__MODULE__)) && + f.args[2] == QuoteNode(:run!) +end + +function _models_process_symbol(expr) + if Meta.isexpr(expr, :.) && length(expr.args) == 2 && expr.args[1] == :models && expr.args[2] isa QuoteNode + return (:Default, expr.args[2].value) + end + return nothing +end + +function _strip_return_expr(expr) + if Meta.isexpr(expr, :return) + return length(expr.args) == 0 ? nothing : expr.args[1] + elseif Meta.isexpr(expr, :block) + args = Any[] + for (i, arg) in enumerate(expr.args) + if Meta.isexpr(arg, :return) + i == length(expr.args) || error("Cannot inline model body with a non-final top-level `return`.") + stripped = _strip_return_expr(arg) + !isnothing(stripped) && push!(args, stripped) + else + _contains_return(arg) && error("Cannot inline model body with a nested `return`.") + push!(args, arg) + end + end + return Expr(:block, args...) + end + return expr +end + +function _contains_return(expr) + Meta.isexpr(expr, :return) && return true + expr isa Expr || return false + return any(_contains_return, expr.args) +end + +function _block_statements(expr) + isnothing(expr) && return Any[] + Meta.isexpr(expr, :block) && return Any[expr.args...] + return Any[expr] +end + +function _sprint_expr(expr) + io = IOBuffer() + Base.show_unquoted(io, expr, 0, 0) + return String(take!(io)) +end diff --git a/src/mtg/GraphSimulation.jl b/src/mtg/GraphSimulation.jl index c53e43eef..8a2595527 100644 --- a/src/mtg/GraphSimulation.jl +++ b/src/mtg/GraphSimulation.jl @@ -29,7 +29,7 @@ struct GraphSimulation{T,S,U,O,V,TS,MS} models::Dict{Symbol,U} model_specs::MS outputs::Dict{Symbol,O} - outputs_index::Dict{Symbol, Int} + outputs_index::Dict{Symbol,Int} temporal_state::TS is_multirate::Bool end @@ -103,7 +103,7 @@ convert_outputs(out, DataFrames) # Another, possibly better way would be to just create the DataFrame directly from the outputs # and then remove the RefVector columns and replace the node one, hmm function convert_outputs(outs::Dict{Symbol,O} where O, sink; refvectors=false, no_value=nothing) - ret = Dict{Symbol, sink}() + ret = Dict{Symbol,sink}() for (organ, status_vector) in outs # remove RefVector variables refv = () @@ -120,7 +120,7 @@ function convert_outputs(outs::Dict{Symbol,O} where O, sink; refvectors=false, n @warn "No instance found at the $organ scale, no output available, removing it from the Dict" continue end - + # Get the new NamedTuple type refv_nt = NamedTuple{refv} @@ -128,12 +128,12 @@ function convert_outputs(outs::Dict{Symbol,O} where O, sink; refvectors=false, n vector_named_tuple_1 = NamedTuple(status_vector[1]) # replace the MTG node var with the id (MTG nodes aren't CSV-friendly) - filtered_named_tuple = (;node=MultiScaleTreeGraph.node_id(vector_named_tuple_1.node),Base.structdiff(vector_named_tuple_1, refv_nt)...) + filtered_named_tuple = (; node=MultiScaleTreeGraph.node_id(vector_named_tuple_1.node), Base.structdiff(vector_named_tuple_1, refv_nt)...) filtered_vector_named_tuple = Vector{typeof(filtered_named_tuple)}(undef, length(status_vector)) for i in 1:length(status_vector) vector_named_tuple_i = NamedTuple(status_vector[i]) - filtered_vector_named_tuple[i] = (;node=MultiScaleTreeGraph.node_id(vector_named_tuple_i.node), Base.structdiff(vector_named_tuple_i, refv_nt)...) + filtered_vector_named_tuple[i] = (; node=MultiScaleTreeGraph.node_id(vector_named_tuple_i.node), Base.structdiff(vector_named_tuple_i, refv_nt)...) end ret[organ] = sink(filtered_vector_named_tuple) @@ -142,16 +142,16 @@ function convert_outputs(outs::Dict{Symbol,O} where O, sink; refvectors=false, n end # TODO adapt these to new output structure or remove them -function outputs(outs::Dict{Symbol, O} where O, key::Symbol) +function outputs(outs::Dict{Symbol,O} where O, key::Symbol) Tables.columns(convert_outputs(outs, Vector{NamedTuple}))[key] end -function outputs(outs::Dict{Symbol, O} where O, i::T) where {T<:Integer} +function outputs(outs::Dict{Symbol,O} where O, i::T) where {T<:Integer} Tables.columns(convert_outputs(outs, Vector{NamedTuple}))[i] end # ModelLists now return outputs as a TimeStepTable{Status}, conversion is straightforward function convert_outputs(out::TimeStepTable{T} where T, sink) - @assert Tables.istable(sink) "The sink argument must be compatible with the Tables.jl interface (`Tables.istable(sink)` must return `true`, *e.g.* `DataFrame`)" + @assert Tables.istable(sink) "The sink argument must be compatible with the Tables.jl interface (`Tables.istable(sink)` must return `true`, *e.g.* `DataFrame`)" return sink(out) end diff --git a/test/runtests.jl b/test/runtests.jl index 99cc88cdb..bdece1116 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -58,6 +58,10 @@ include("helper-functions.jl") include("test-simulation.jl") end + @testset "Compiled model source" begin + include("test-compiled-model.jl") + end + @testset "Statistics" begin include("test-statistics.jl") end diff --git a/test/test-compiled-model.jl b/test/test-compiled-model.jl new file mode 100644 index 000000000..e7c2f2168 --- /dev/null +++ b/test/test-compiled-model.jl @@ -0,0 +1,410 @@ +@testset "Merged model source compiler" begin + mapping = ModelMapping( + process4=Process4Model(), + process3=Process3Model(), + process2=Process2Model(), + process1=Process1Model(1.0); + status=(var0=1.0,) + ) + + script = compile_model(mapping; function_name=:compiled_dummy_model!) + @test occursin("function compiled_dummy_model!(models, status, meteo, constants, extra)", script) + @test occursin("using PlantSimEngine", script) + @test occursin("using PlantSimEngine.Examples", script) + @test occursin("model: Process4Model | process: process4 | scale: Default", script) + @test occursin("inputs: var0 | outputs: var1, var2", script) + @test occursin("process4_model = models.process4", script) + @test occursin("source: " * joinpath(pkgdir(PlantSimEngine), "examples", "dummy.jl"), script) + @test occursin("method: run!(::Process1Model", script) + @test occursin("process4_status.var1 = process4_status.var0 + 0.01", script) + @test occursin("process1_status.var3 = process1_model.a + process1_status.var1 * process1_status.var2", script) + @test !occursin("PlantSimEngine.run!(models.process", script) + + fast_script = compile_model(mapping; function_name=:compiled_dummy_model_fast!, mode=:fast) + @test occursin("status.var1 = status.var0 + 0.01", fast_script) + @test !occursin("process4_model = models.process4", fast_script) + @test_throws "Unsupported compiled model mode" compile_model(mapping; mode=:compact) + + module_name = gensym(:CompiledModelTest) + compiled_mod = Module(module_name) + script_path = tempname() * ".jl" + write(script_path, script) + Base.include(compiled_mod, script_path) + + model_list = PlantSimEngine._modellist_from_model_mapping(mapping) + compiled_status = deepcopy(status(model_list)) + meteo = Atmosphere(T=20.0, Wind=1.0, P=101.3, Rh=0.65) + Core.eval(compiled_mod, :compiled_dummy_model!)(model_list.models, compiled_status, meteo, PlantMeteo.Constants(), nothing) + + normal_mapping = copy(mapping) + run!(normal_mapping, meteo; executor=SequentialEx()) + normal_status = status(PlantSimEngine._modellist_from_model_mapping(normal_mapping)) + + @test compiled_status.var1 == normal_status.var1 + @test compiled_status.var2 == normal_status.var2 + @test compiled_status.var3 == normal_status.var3 + @test compiled_status.var4 == normal_status.var4 + @test compiled_status.var5 == normal_status.var5 + @test compiled_status.var6 == normal_status.var6 + + path = tempname() * ".jl" + @test write_compiled_model(path, mapping; function_name=:compiled_dummy_model!) == path + @test read(path, String) == script +end + +PlantSimEngine.@process "compiled_default_child" verbose = false +struct CompiledDefaultChildModel <: AbstractCompiled_Default_ChildModel end +PlantSimEngine.inputs_(::CompiledDefaultChildModel) = (x=-Inf,) +PlantSimEngine.outputs_(::CompiledDefaultChildModel) = (y=-Inf,) +function PlantSimEngine.run!(::CompiledDefaultChildModel, models, status, meteo, constants=nothing, extra=nothing) + begin + tmp = status.x + 2.0 + status.y = tmp + end +end + +PlantSimEngine.@process "compiled_default_parent" verbose = false +struct CompiledDefaultParentModel <: AbstractCompiled_Default_ParentModel end +PlantSimEngine.dep(::CompiledDefaultParentModel) = (compiled_default_child=AbstractCompiled_Default_ChildModel,) +PlantSimEngine.inputs_(::CompiledDefaultParentModel) = (x=-Inf,) +PlantSimEngine.outputs_(::CompiledDefaultParentModel) = (z=-Inf,) +function PlantSimEngine.run!(::CompiledDefaultParentModel, models, status, meteo, constants=nothing, extra=nothing) + let offset = 3.0 + child_model = getproperty(models, :compiled_default_child) + child_runner = PlantSimEngine.run! + child_runner(child_model, models, status, meteo) + status.z = status.y + offset + end +end + +@testset "Merged single-scale compiler handles default args and nested blocks" begin + mapping = ModelMapping( + compiled_default_parent=CompiledDefaultParentModel(), + compiled_default_child=CompiledDefaultChildModel(); + status=(x=5.0,) + ) + + script = compile_model(mapping; function_name=:compiled_default_model!) + @test occursin("tmp = compiled_default_child_status.x + 2.0", script) + @test occursin("compiled_default_child_status.y = tmp", script) + @test !occursin("run!(models.compiled_default_child", script) + + compiled_mod = Module(gensym(:CompiledDefaultShapeTest)) + script_path = tempname() * ".jl" + write(script_path, script) + Base.include(compiled_mod, script_path) + + model_list = PlantSimEngine._modellist_from_model_mapping(mapping) + compiled_status = deepcopy(status(model_list)) + meteo = Atmosphere(T=20.0, Wind=1.0, Rh=0.65) + Core.eval(compiled_mod, :compiled_default_model!)(model_list.models, compiled_status, meteo, PlantMeteo.Constants(), nothing) + + @test compiled_status.y == 7.0 + @test compiled_status.z == 10.0 +end + +PlantSimEngine.@process "compiled_alias_child" verbose = false +struct CompiledAliasChildModel <: AbstractCompiled_Alias_ChildModel end +PlantSimEngine.inputs_(::CompiledAliasChildModel) = (x=-Inf,) +PlantSimEngine.outputs_(::CompiledAliasChildModel) = (y=-Inf,) +function PlantSimEngine.run!(::CompiledAliasChildModel, models, status, meteo, constants=nothing, extra=nothing) + status.y = status.x + 1.0 +end + +PlantSimEngine.@process "compiled_alias_parent" verbose = false +struct CompiledAliasParentModel <: AbstractCompiled_Alias_ParentModel end +PlantSimEngine.dep(::CompiledAliasParentModel) = (compiled_alias_child=AbstractCompiled_Alias_ChildModel,) +PlantSimEngine.inputs_(::CompiledAliasParentModel) = (x=-Inf,) +PlantSimEngine.outputs_(::CompiledAliasParentModel) = (z=-Inf,) +function PlantSimEngine.run!(::CompiledAliasParentModel, models, status, meteo, constants=nothing, extra=nothing) + dep_process = Symbol("compiled_alias_child") + run!(getfield(models, dep_process), models, status, meteo, constants, extra) + status.z = status.y + 1.0 +end + +@testset "Merged compiler rejects unsupported hard dependency call shapes" begin + mapping = ModelMapping( + compiled_alias_parent=CompiledAliasParentModel(), + compiled_alias_child=CompiledAliasChildModel(); + status=(x=1.0,) + ) + + @test_throws "Failed to compile process `compiled_alias_parent`" compile_model(mapping; function_name=:compiled_alias_model!) +end + +function compiled_model_dynamic_fixture() + mtg = import_mtg_example() + meteo = Weather([ + Atmosphere(T=20.0, Wind=1.0, Rh=0.65), + Atmosphere(T=25.0, Wind=0.5, Rh=0.8), + ]) + + mapping = ModelMapping( + :Scene => ToyDegreeDaysCumulModel(), + :Plant => ( + MultiScaleModel( + model=ToyLAIModel(), + mapped_variables=[:TT_cu => (:Scene => :TT_cu)], + ), + Beer(0.6), + MultiScaleModel( + model=ToyCAllocationModel(), + mapped_variables=[ + :carbon_assimilation => [:Leaf], + :carbon_demand => [:Leaf, :Internode], + :carbon_allocation => [:Leaf, :Internode], + ], + ), + MultiScaleModel( + model=ToyPlantRmModel(), + mapped_variables=[:Rm_organs => [:Leaf => :Rm, :Internode => :Rm]], + ), + ), + :Internode => ( + MultiScaleModel( + model=ToyCDemandModel(optimal_biomass=10.0, development_duration=200.0), + mapped_variables=[:TT => (:Scene => :TT)], + ), + MultiScaleModel( + model=ToyInternodeEmergence(TT_emergence=20.0), + mapped_variables=[:TT_cu => (:Scene => :TT_cu)], + ), + ToyMaintenanceRespirationModel(1.5, 0.06, 25.0, 0.6, 0.004), + Status(carbon_biomass=1.0), + ), + :Leaf => ( + MultiScaleModel( + model=ToyAssimModel(), + mapped_variables=[:soil_water_content => (:Soil => :soil_water_content), :aPPFD => (:Plant => :aPPFD)], + ), + MultiScaleModel( + model=ToyCDemandModel(optimal_biomass=10.0, development_duration=200.0), + mapped_variables=[:TT => (:Scene => :TT)], + ), + ToyMaintenanceRespirationModel(2.1, 0.06, 25.0, 1.0, 0.025), + Status(carbon_biomass=1.0), + ), + :Soil => (ToySoilWaterModel(),), + ) + + outputs = Dict( + :Leaf => (:carbon_assimilation, :carbon_demand, :soil_water_content, :carbon_allocation), + :Internode => (:carbon_allocation, :TT_cu_emergence), + :Plant => (:carbon_allocation,), + :Soil => (:soil_water_content,), + ) + + return mtg, meteo, mapping, outputs +end + +@testset "Merged multiscale model source compiler" begin + mtg, meteo, mapping, outputs = compiled_model_dynamic_fixture() + sim = PlantSimEngine.GraphSimulation(mtg, mapping, nsteps=PlantSimEngine.get_nsteps(meteo), check=true, outputs=outputs) + script = compile_model(sim; function_name=:compiled_mtg_model!) + + @test occursin("function compiled_mtg_model!(sim, meteo, constants)", script) + @test occursin("Mode: same-rate multiscale MTG", script) + @test occursin("variables are scale-scoped through each Status", script) + @test occursin("PlantSimEngine._assert_compiled_sim_compatible", script) + @test occursin("Model compatibility signature", script) + @test occursin("idx_limit = length(statuses[:Leaf])", script) + @test occursin("while idx <= idx_limit", script) + @test occursin("scale: Leaf | initial_statuses: 2", script) + @test occursin("model: ToyCAllocationModel | process: carbon_allocation | scale: Plant", script) + + compiled_mod = Module(gensym(:CompiledMTGModelTest)) + script_path = tempname() * ".jl" + write(script_path, script) + Base.include(compiled_mod, script_path) + compiled_outputs = Core.eval(compiled_mod, :compiled_mtg_model!)(sim, meteo, PlantMeteo.Constants()) + + mtg_normal, meteo_normal, mapping_normal, outputs_normal = compiled_model_dynamic_fixture() + sim_normal = PlantSimEngine.GraphSimulation(mtg_normal, mapping_normal, nsteps=PlantSimEngine.get_nsteps(meteo_normal), check=true, outputs=outputs_normal) + normal_outputs = run!(sim_normal, meteo_normal; executor=SequentialEx()) + + @test length(mtg) == length(mtg_normal) == 9 + @test length(sim.statuses[:Scene]) == length(sim_normal.statuses[:Scene]) == 1 + @test length(sim.statuses[:Soil]) == length(sim_normal.statuses[:Soil]) == 1 + @test length(sim.statuses[:Plant]) == length(sim_normal.statuses[:Plant]) == 1 + @test length(sim.statuses[:Internode]) == length(sim_normal.statuses[:Internode]) == 3 + @test length(sim.statuses[:Leaf]) == length(sim_normal.statuses[:Leaf]) == 3 + + @test sim.statuses[:Leaf][1].carbon_assimilation == sim_normal.statuses[:Leaf][1].carbon_assimilation + @test sim.statuses[:Leaf][2].carbon_assimilation == sim_normal.statuses[:Leaf][2].carbon_assimilation + @test sim.statuses[:Plant][1].carbon_assimilation == sim_normal.statuses[:Plant][1].carbon_assimilation + @test sim.statuses[:Plant][1].carbon_assimilation isa PlantSimEngine.RefVector + @test getfield(sim.statuses[:Plant][1].carbon_assimilation, :v)[1] === PlantSimEngine.refvalue(sim.statuses[:Leaf][1], :carbon_assimilation) + + @test convert_outputs(compiled_outputs, DataFrames.DataFrame) == convert_outputs(normal_outputs, DataFrames.DataFrame) + + mtg_wrong = import_mtg_example() + meteo_wrong = Atmosphere(T=20.0, Wind=1.0, Rh=0.65) + mapping_wrong = ModelMapping(:Soil => (ToySoilWaterModel(),)) + sim_wrong = PlantSimEngine.GraphSimulation(mtg_wrong, mapping_wrong, nsteps=1, check=true, outputs=Dict(:Soil => (:soil_water_content,))) + @test_throws "Compiled model was generated for a different model mapping" Core.eval(compiled_mod, :compiled_mtg_model!)( + sim_wrong, + meteo_wrong, + PlantMeteo.Constants(), + ) +end + +function compiled_model_multirate_fixture() + mtg = MultiScaleTreeGraph.Node(MultiScaleTreeGraph.NodeMTG("/", :Scene, 1, 0)) + plant = MultiScaleTreeGraph.Node(mtg, MultiScaleTreeGraph.NodeMTG("+", :Plant, 1, 1)) + MultiScaleTreeGraph.Node(mtg, MultiScaleTreeGraph.NodeMTG("+", :Soil, 1, 1)) + internode = MultiScaleTreeGraph.Node(plant, MultiScaleTreeGraph.NodeMTG("/", :Internode, 1, 2)) + MultiScaleTreeGraph.Node(internode, MultiScaleTreeGraph.NodeMTG("+", :Leaf, 1, 2)) + MultiScaleTreeGraph.Node(internode, MultiScaleTreeGraph.NodeMTG("+", :Leaf, 1, 2)) + + daily = ClockSpec(24.0, 1.0) + hourly = 1.0 + mapping = ModelMapping( + :Scene => ( + ModelSpec(ToyDegreeDaysCumulModel()) |> TimeStepModel(daily), + ), + :Plant => ( + ModelSpec(ToyLAIModel()) |> + MultiScaleModel([:TT_cu => (:Scene => :TT_cu)]) |> + TimeStepModel(daily), + ModelSpec(Beer(0.6)) |> TimeStepModel(hourly), + ModelSpec(ToyCAllocationModel()) |> + MultiScaleModel([ + :carbon_assimilation => [:Leaf], + :carbon_demand => [:Leaf, :Internode], + :carbon_allocation => [:Leaf, :Internode], + ]) |> + InputBindings(; carbon_assimilation=(process=process(ToyAssimModel()), var=:carbon_assimilation, scale=:Leaf, policy=Integrate())) |> + TimeStepModel(daily), + ModelSpec(ToyPlantRmModel()) |> + MultiScaleModel([:Rm_organs => [:Leaf => :Rm, :Internode => :Rm]]) |> + TimeStepModel(daily), + ), + :Internode => ( + ModelSpec(ToyCDemandModel(optimal_biomass=10.0, development_duration=200.0)) |> + MultiScaleModel([:TT => (:Scene => :TT)]) |> + TimeStepModel(daily), + ModelSpec(ToyInternodeEmergence(TT_emergence=1.0e6)) |> + MultiScaleModel([:TT_cu => (:Scene => :TT_cu)]) |> + TimeStepModel(daily), + ModelSpec(ToyMaintenanceRespirationModel(1.5, 0.06, 25.0, 0.6, 0.004)) |> TimeStepModel(daily), + Status(carbon_biomass=1.0), + ), + :Leaf => ( + ModelSpec(ToyAssimModel()) |> + MultiScaleModel([:soil_water_content => (:Soil => :soil_water_content), :aPPFD => (:Plant => :aPPFD)]) |> + TimeStepModel(hourly), + ModelSpec(ToyCDemandModel(optimal_biomass=10.0, development_duration=200.0)) |> + MultiScaleModel([:TT => (:Scene => :TT)]) |> + TimeStepModel(daily), + ModelSpec(ToyMaintenanceRespirationModel(2.1, 0.06, 25.0, 1.0, 0.025)) |> TimeStepModel(daily), + Status(carbon_biomass=1.0), + ), + :Soil => ( + ModelSpec(ToySoilWaterModel()) |> TimeStepModel(daily), + ), + ) + + outputs = Dict( + :Leaf => (:carbon_assimilation, :aPPFD, :carbon_demand), + :Plant => (:LAI, :carbon_offer, :Rm), + :Scene => (:TT, :TT_cu), + :Soil => (:soil_water_content,), + :Internode => (:carbon_demand, :TT_cu_emergence), + ) + meteo = Weather(repeat([Atmosphere(T=20.0, Wind=1.0, Rh=0.65, Ri_PAR_f=300.0)], 26)) + + return mtg, meteo, mapping, outputs +end + +@testset "Merged multiscale model compiler supports multirate" begin + mtg, meteo, mapping, outputs = compiled_model_multirate_fixture() + sim = PlantSimEngine.GraphSimulation(mtg, mapping, nsteps=PlantSimEngine.get_nsteps(meteo), check=true, outputs=outputs) + script = compile_model(sim; function_name=:compiled_multirate_mtg_model!) + + @test occursin("Mode: multi-rate multiscale MTG", script) + @test occursin("tracked_outputs = nothing", script) + @test occursin("PlantSimEngine.resolve_inputs_from_temporal_state!", script) + @test occursin("PlantSimEngine.update_temporal_state_outputs!", script) + @test occursin("PlantSimEngine.update_requested_outputs!", script) + @test occursin("_mr_Leaf_carbon_assimilation_model = (models_by_scale[:Leaf]).carbon_assimilation", script) + @test occursin("_mr_Leaf_carbon_assimilation_model_clock = PlantSimEngine._model_clock", script) + @test occursin("idx_limit = length(statuses[:Leaf])", script) + @test occursin("while idx <= idx_limit", script) + + compiled_mod = Module(gensym(:CompiledMultirateMTGModelTest)) + script_path = tempname() * ".jl" + write(script_path, script) + Base.include(compiled_mod, script_path) + tracked = OutputRequest(:Leaf, :carbon_assimilation; name=:leaf_assimilation, process=process(ToyAssimModel())) + compiled_outputs, compiled_requested = Core.eval(compiled_mod, :compiled_multirate_mtg_model!)( + sim, + meteo, + PlantMeteo.Constants(); + tracked_outputs=tracked, + return_requested_outputs=true, + requested_outputs_sink=DataFrames.DataFrame, + ) + + mtg_normal, meteo_normal, mapping_normal, outputs_normal = compiled_model_multirate_fixture() + sim_normal = PlantSimEngine.GraphSimulation(mtg_normal, mapping_normal, nsteps=PlantSimEngine.get_nsteps(meteo_normal), check=true, outputs=outputs_normal) + normal_outputs, normal_requested = run!( + sim_normal, + meteo_normal; + executor=SequentialEx(), + tracked_outputs=tracked, + return_requested_outputs=true, + requested_outputs_sink=DataFrames.DataFrame, + ) + + @test sim.temporal_state.last_run == sim_normal.temporal_state.last_run + @test convert_outputs(compiled_outputs, DataFrames.DataFrame) == convert_outputs(normal_outputs, DataFrames.DataFrame) + @test compiled_requested[:leaf_assimilation] == normal_requested[:leaf_assimilation] +end + +PlantSimEngine.@process "compiled_child" verbose = false +struct CompiledChildModel <: AbstractCompiled_ChildModel end +PlantSimEngine.inputs_(::CompiledChildModel) = (x=-Inf,) +PlantSimEngine.outputs_(::CompiledChildModel) = (y=-Inf,) +function PlantSimEngine.run!(::CompiledChildModel, models, status, meteo, constants=nothing, extra=nothing) + status.y = status.x + 1.0 +end + +PlantSimEngine.@process "compiled_parent" verbose = false +struct CompiledParentModel <: AbstractCompiled_ParentModel end +PlantSimEngine.inputs_(::CompiledParentModel) = (x=-Inf,) +PlantSimEngine.outputs_(::CompiledParentModel) = (total=-Inf,) +PlantSimEngine.dep(::CompiledParentModel) = (compiled_child=AbstractCompiled_ChildModel => (:CompiledChildScale,),) +function PlantSimEngine.run!(::CompiledParentModel, models, status, meteo, constants=nothing, sim_object=nothing) + child_status = sim_object.statuses[:CompiledChildScale][1] + run!(sim_object.models[:CompiledChildScale].compiled_child, models, child_status, meteo, constants) + status.total = status.x + child_status.y +end + +@testset "Merged multiscale compiler inlines cross-scale hard dependencies" begin + mtg = MultiScaleTreeGraph.Node(MultiScaleTreeGraph.NodeMTG("/", :CompiledParentScale, 1, 0)) + MultiScaleTreeGraph.Node(mtg, MultiScaleTreeGraph.NodeMTG("/", :CompiledChildScale, 1, 1)) + mapping = ModelMapping( + :CompiledParentScale => (CompiledParentModel(), Status(x=10.0)), + :CompiledChildScale => (CompiledChildModel(), Status(x=2.0)), + ) + sim = PlantSimEngine.GraphSimulation( + mtg, + mapping, + nsteps=1, + check=true, + outputs=Dict(:CompiledParentScale => (:total,), :CompiledChildScale => (:y,)) + ) + + script = compile_model(sim; function_name=:compiled_cross_scale_hard_dep!) + @test occursin("sim.models[:CompiledChildScale]", script) + @test !occursin("run!(sim.models[:CompiledChildScale].compiled_child", script) + + compiled_mod = Module(gensym(:CompiledCrossScaleHardDepTest)) + script_path = tempname() * ".jl" + write(script_path, script) + Base.include(compiled_mod, script_path) + Core.eval(compiled_mod, :compiled_cross_scale_hard_dep!)(sim, Atmosphere(T=20.0, Wind=1.0, Rh=0.65), PlantMeteo.Constants()) + + @test sim.statuses[:CompiledChildScale][1].y == 3.0 + @test sim.statuses[:CompiledParentScale][1].total == 13.0 +end