From 33cc0afa65fcfa8307a4e881d3098b52c2fe2e17 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?R=C3=A9mi=20Vezy?= <VEZY@users.noreply.github.com>
Date: Mon, 22 Jul 2024 18:15:02 +0200
Subject: [PATCH] Output variables for unexisting node types now OK (can appear
 after in the simulation)

---
 src/mtg/initialisation.jl |  2 +-
 src/mtg/save_results.jl   | 25 ++++++++++++++++++-------
 2 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/src/mtg/initialisation.jl b/src/mtg/initialisation.jl
index beaa641..173748d 100644
--- a/src/mtg/initialisation.jl
+++ b/src/mtg/initialisation.jl
@@ -318,7 +318,7 @@ function init_simulation(mtg, mapping; nsteps=1, outputs=nothing, type_promotion
 
     models = Dict(first(m) => parse_models(get_models(last(m))) for m in mapping)
 
-    outputs = pre_allocate_outputs(statuses, outputs, nsteps, check=check)
+    outputs = pre_allocate_outputs(statuses, status_templates, reverse_multiscale_mapping, vars_need_init, outputs, nsteps, type_promotion=type_promotion, check=check)
 
     return (; mtg, statuses, status_templates, reverse_multiscale_mapping, vars_need_init, dependency_graph=dep(mapping, verbose=verbose), models, outputs)
 end
\ No newline at end of file
diff --git a/src/mtg/save_results.jl b/src/mtg/save_results.jl
index 1897c87..a621b7c 100644
--- a/src/mtg/save_results.jl
+++ b/src/mtg/save_results.jl
@@ -109,9 +109,9 @@ julia> collect(keys(preallocated_vars["Leaf"]))
  :carbon_demand
 ```
 """
-function pre_allocate_outputs(statuses, outs, nsteps; check=true)
+function pre_allocate_outputs(statuses, statuses_template, reverse_multiscale_mapping, vars_need_init, outs, nsteps; type_promotion=nothing, check=true)
     outs_ = copy(outs)
-
+    statuses_ = copy(statuses)
     # Checking that organs in outputs exist in the mtg (in the statuses):
     if !all(i in keys(statuses) for i in keys(outs_))
         not_in_statuses = setdiff(keys(outs_), keys(statuses))
@@ -135,11 +135,22 @@ function pre_allocate_outputs(statuses, outs, nsteps; check=true)
     for (organ, vars) in outs_ # organ = "Leaf"; vars = outs_[organ]
         if length(statuses[organ]) == 0
             # The organ is not found in the mtg, we return an info and get along (it might be created during the simulation):
-            check && @info "You required outputs for organ $organ, but this organ is not found in the provided MTG."
-            continue
+            check && @info "You required outputs for organ $organ, but this organ is not found in the provided MTG at this point."
+
+            status_from_template = PlantSimEngine.init_node_status!(
+                MultiScaleTreeGraph.Node(MultiScaleTreeGraph.NodeMTG("/", organ, 0, 0)),
+                statuses,
+                statuses_template,
+                reverse_multiscale_mapping,
+                vars_need_init,
+                type_promotion;
+                check=check
+            )
+
+            statuses_[organ] = [status_from_template]
         end
-        if !all(i in collect(keys(statuses[organ][1])) for i in vars)
-            not_in_statuses = (setdiff(vars, keys(statuses[organ][1]))...,)
+        if !all(i in collect(keys(statuses_[organ][1])) for i in vars)
+            not_in_statuses = (setdiff(vars, keys(statuses_[organ][1]))...,)
             plural = length(not_in_statuses) == 1 ? "" : "s"
             e = string(
                 "You requested outputs for variable", plural, " ",
@@ -168,7 +179,7 @@ function pre_allocate_outputs(statuses, outs, nsteps; check=true)
     end
 
     # Making the pre-allocated outputs:
-    Dict(organ => Dict(var => [typeof(statuses[organ][1][var])[] for n in 1:nsteps] for var in vars) for (organ, vars) in outs_)
+    Dict(organ => Dict(var => [typeof(statuses_[organ][1][var])[] for n in 1:nsteps] for var in vars) for (organ, vars) in outs_)
     # Note: we use the type of the variable from the first status for each organ to pre-allocate the outputs, because they are
     # all the same type for others.
 end