Fix little issues + attempt at addressing main blocker : disconnect timestep-mapped variables from their source, to avoid overwriting source (inputs to non-default-timestep models are changed by the accumulation function so can't be a simple Ref to source)

Author: Samuel-amap

src/dependencies/dependency_graph.jl

@@ -89,82 +89,4 @@ function Base.show(io::IO, ::MIME"text/plain", t::DependencyGraph)
     else
         draw_dependency_graph(io, t)
     end
-end
-
-"""
-    variables_multiscale(node, organ, mapping, st=NamedTuple())
-
-Get the variables of a HardDependencyNode, taking into account the multiscale mapping, *i.e.*
-defining variables as `MappedVar` if they are mapped to another scale. The default values are 
-taken from the model if not given by the user (`st`), and are marked as `UninitializedVar` if 
-they are inputs of the node.
-
-Return a NamedTuple with the variables and their default values.
-
-# Arguments
-
-- `node::HardDependencyNode`: the node to get the variables from.
-- `organ::String`: the organ type, *e.g.* "Leaf".
-- `vars_mapping::Dict{String,T}`: the mapping of the models (see details below).
-- `st::NamedTuple`: an optional named tuple with default values for the variables.
-
-# Details
-
-The `vars_mapping` is a dictionary with the organ type as key and a dictionary as value. It is 
-computed from the user mapping like so:
-"""
-function variables_multiscale(node, organ, vars_mapping, st=NamedTuple())
-    node_vars = variables(node) # e.g. (inputs = (:var1=-Inf, :var2=-Inf), outputs = (:var3=-Inf,))
-    ins = node_vars.inputs
-    ins_variables = keys(ins)
-    outs_variables = keys(node_vars.outputs)
-    defaults = merge(node_vars...)
-    map((inputs=ins_variables, outputs=outs_variables)) do vars # Map over vars from :inputs and vars from :outputs
-        vars_ = Vector{Pair{Symbol,Any}}()
-        for var in vars # e.g. var = :carbon_biomass
-            if var in keys(st)
-                #If the user has given a status, we use it as default value.
-                default = st[var]
-            elseif var in ins_variables
-                # Otherwise, we use the default value given by the model:
-                # If the variable is an input, we mark it as uninitialized:
-                default = UninitializedVar(var, defaults[var])
-            else
-                # If the variable is an output, we use the default value given by the model:
-                default = defaults[var]
-            end
-
-            if haskey(vars_mapping[organ], var)
-                organ_mapped, organ_mapped_var = _node_mapping(vars_mapping[organ][var])
-                push!(vars_, var => MappedVar(organ_mapped, var, organ_mapped_var, default))
-                #* We still check if the variable also exists wrapped in PreviousTimeStep, because one model could use the current 
-                #* values, and another one the previous values.
-                if haskey(vars_mapping[organ], PreviousTimeStep(var, node.process))
-                    organ_mapped, organ_mapped_var = _node_mapping(vars_mapping[organ][PreviousTimeStep(var, node.process)])
-                    push!(vars_, var => MappedVar(organ_mapped, PreviousTimeStep(var, node.process), organ_mapped_var, default))
-                end
-            elseif haskey(vars_mapping[organ], PreviousTimeStep(var, node.process))
-                # If not found in the current time step, we check if the variable is mapped to the previous time step:
-                organ_mapped, organ_mapped_var = _node_mapping(vars_mapping[organ][PreviousTimeStep(var, node.process)])
-                push!(vars_, var => MappedVar(organ_mapped, PreviousTimeStep(var, node.process), organ_mapped_var, default))
-            else
-                # Else we take the default value:
-                push!(vars_, var => default)
-            end
-        end
-        return (; vars_...,)
-    end
-end
-
-function _node_mapping(var_mapping::Pair{String,Symbol})
-    # One organ is mapped to the variable:
-    return SingleNodeMapping(first(var_mapping)), last(var_mapping)
-end
-
-function _node_mapping(var_mapping)
-    # Several organs are mapped to the variable:
-    organ_mapped = MultiNodeMapping([first(i) for i in var_mapping])
-    organ_mapped_var = [last(i) for i in var_mapping]
-
-    return organ_mapped, organ_mapped_var
 end

src/dependencies/hard_dependencies.jl

@@ -110,6 +110,92 @@ function initialise_all_as_hard_dependency_node(models, scale)
     return dep_graph
 end
 
+# Samuel : this requires the orchestrator, which requires the dependency graph
+# Leaving it in dependency_graph.jl causes forward declaration issues, moving it here as a quick protoyping hack, it might not be the ideal spot
+"""
+    variables_multiscale(node, organ, mapping, st=NamedTuple())
+
+Get the variables of a HardDependencyNode, taking into account the multiscale mapping, *i.e.*
+defining variables as `MappedVar` if they are mapped to another scale. The default values are 
+taken from the model if not given by the user (`st`), and are marked as `UninitializedVar` if 
+they are inputs of the node.
+
+Return a NamedTuple with the variables and their default values.
+
+# Arguments
+
+- `node::HardDependencyNode`: the node to get the variables from.
+- `organ::String`: the organ type, *e.g.* "Leaf".
+- `vars_mapping::Dict{String,T}`: the mapping of the models (see details below).
+- `st::NamedTuple`: an optional named tuple with default values for the variables.
+
+# Details
+
+The `vars_mapping` is a dictionary with the organ type as key and a dictionary as value. It is 
+computed from the user mapping like so:
+"""
+function variables_multiscale(node, organ, vars_mapping, st=NamedTuple(), orchestrator::Orchestrator2=Orchestrator2())
+    node_vars = variables(node) # e.g. (inputs = (:var1=-Inf, :var2=-Inf), outputs = (:var3=-Inf,))
+    ins = node_vars.inputs
+    ins_variables = keys(ins)
+    outs_variables = keys(node_vars.outputs)
+    defaults = merge(node_vars...)
+    map((inputs=ins_variables, outputs=outs_variables)) do vars # Map over vars from :inputs and vars from :outputs
+        vars_ = Vector{Pair{Symbol,Any}}()
+        for var in vars # e.g. var = :carbon_biomass
+            if var in keys(st)
+                #If the user has given a status, we use it as default value.
+                default = st[var]
+            elseif var in ins_variables
+                # Otherwise, we use the default value given by the model:
+                # If the variable is an input, we mark it as uninitialized:
+                default = UninitializedVar(var, defaults[var])
+            else
+                # If the variable is an output, we use the default value given by the model:
+                default = defaults[var]
+            end
+
+            # TODO no idea how this meshes with refvector situations or previoustimestep
+            if is_timestep_mapped(organ => var, orchestrator, search_inputs_only=true)
+                
+                push!(vars_, var => default)
+            else
+
+                if haskey(vars_mapping[organ], var)
+                    organ_mapped, organ_mapped_var = _node_mapping(vars_mapping[organ][var])
+                    push!(vars_, var => MappedVar(organ_mapped, var, organ_mapped_var, default))
+                    #* We still check if the variable also exists wrapped in PreviousTimeStep, because one model could use the current 
+                    #* values, and another one the previous values.
+                    if haskey(vars_mapping[organ], PreviousTimeStep(var, node.process))
+                        organ_mapped, organ_mapped_var = _node_mapping(vars_mapping[organ][PreviousTimeStep(var, node.process)])
+                        push!(vars_, var => MappedVar(organ_mapped, PreviousTimeStep(var, node.process), organ_mapped_var, default))
+                    end
+                elseif haskey(vars_mapping[organ], PreviousTimeStep(var, node.process))
+                    # If not found in the current time step, we check if the variable is mapped to the previous time step:
+                    organ_mapped, organ_mapped_var = _node_mapping(vars_mapping[organ][PreviousTimeStep(var, node.process)])
+                    push!(vars_, var => MappedVar(organ_mapped, PreviousTimeStep(var, node.process), organ_mapped_var, default))
+                else
+                    # Else we take the default value:
+                    push!(vars_, var => default)
+                end
+            end
+        end
+        return (; vars_...,)
+    end
+end
+
+function _node_mapping(var_mapping::Pair{String,Symbol})
+    # One organ is mapped to the variable:
+    return SingleNodeMapping(first(var_mapping)), last(var_mapping)
+end
+
+function _node_mapping(var_mapping)
+    # Several organs are mapped to the variable:
+    organ_mapped = MultiNodeMapping([first(i) for i in var_mapping])
+    organ_mapped_var = [last(i) for i in var_mapping]
+
+    return organ_mapped, organ_mapped_var
+end
 
 # When we use a mapping (multiscale), we return the set of soft-dependencies (we put the hard-dependencies as their children):
 function hard_dependencies(mapping::Dict{String,T}; verbose::Bool=true, orchestrator::Orchestrator2=Orchestrator2()) where {T}
@@ -148,7 +234,7 @@ function hard_dependencies(mapping::Dict{String,T}; verbose::Bool=true, orchestr
         status_scale = Dict{Symbol,Vector{Pair{Symbol,NamedTuple}}}()
         for (procname, node) in hard_deps[organ].roots # procname = :leaf_surface ; node = hard_deps.roots[procname]
             var = Pair{Symbol,NamedTuple}[]
-            traverse_dependency_graph!(node, x -> variables_multiscale(x, organ, full_vars_mapping, st_scale_user), var)
+            traverse_dependency_graph!(node, x -> variables_multiscale(x, organ, full_vars_mapping, st_scale_user, orchestrator), var)
             push!(status_scale, procname => var)
         end
 

src/dependencies/soft_dependencies.jl

@@ -357,7 +357,7 @@ function add_timestep_data_to_node(soft_dependency_node, orchestrator::Orchestra
     # now we can create the mapping
     for mtsm in orchestrator.non_default_timestep_mapping
         if mtsm.scale == soft_dependency_node.scale && (mtsm.model) == typeof(model_(soft_dependency_node.value))
-            for (var_to, var_from) in mtsm.var_to_var        
+            for (var_from, var_to) in mtsm.var_to_var        
                 if !isnothing(soft_dependency_node.parent)
                     parent = nothing
                     variable_mapping = nothing

src/mtg/initialisation.jl

@@ -35,26 +35,31 @@ function init_statuses(mtg, mapping, dependency_graph=first(hard_dependencies(ma
     # Note 3: we do it before `convert_reference_values!` because we need the variables to be MappedVar{MultiNodeMapping} to get the reverse mapping.
 
     # Convert the MappedVar{SelfNodeMapping} or MappedVar{SingleNodeMapping} to RefValues, and MappedVar{MultiNodeMapping} to RefVectors:
-    convert_reference_values!(mapped_vars)
+    convert_reference_values!(mapped_vars, orchestrator)
 
     # Get the variables that are not initialised or computed by other models in the output:
     vars_need_init = Dict(org => filter(x -> isa(last(x), UninitializedVar), vars) |> keys |> collect for (org, vars) in mapped_vars) |>
                      filter(x -> length(last(x)) > 0)
 
     # Filter out variables that are timestep-mapped by the user, 
-    # as those variables are initialized by another model, but are currently flagged as needing initialization
+    # Since we disconnect outputs from the source variable (as values are changed by the accumulation function,
+    # meaning they differ and we can't just Ref point to the source variable)
+    # they will be currently flagged as needing initialization
     # At this stage, data present in the orchestrator is expected to be valid, so we can take it into account
     # A model with a different timestep can still have unitialized vars found in a node, the meteo, or to be initialized by the user
     # in which case it'll be absent from the timestep mapping, but this needs testing
-    #filter_timestep_mapped_variables!(vars_need_init, orchestrator)
+    # TODO, *however*, this isn't the cleanest in its current state, 
+    # there may be some user initialisation issues that are hidden by this approach
+    # Needs to be checked
+    filter_timestep_mapped_variables!(vars_need_init, orchestrator)
 
 
     # Note: these variables may be present in the MTG attributes, we check that below when traversing the MTG.
 
     # We traverse the MTG to initialise the statuses linked to the nodes:
     statuses = Dict(i => Status[] for i in collect(keys(mapped_vars)))
     MultiScaleTreeGraph.traverse!(mtg) do node # e.g.: node = MultiScaleTreeGraph.get_node(mtg, 5)
-        init_node_status!(node, statuses, mapped_vars, reverse_multiscale_mapping, vars_need_init, type_promotion, check=check)
+        init_node_status!(node, statuses, mapped_vars, reverse_multiscale_mapping, vars_need_init, type_promotion, check=check, orchestrator=orchestrator)
     end
 
     return (; statuses, mapped_vars, reverse_multiscale_mapping, vars_need_init)
@@ -153,7 +158,7 @@ function init_node_status!(node, statuses, mapped_vars, reverse_multiscale_mappi
     end
 
     # Make the node status from the template:
-    st = status_from_template(st_template)
+    st = status_from_template(st_template, symbol(node), orchestrator)
 
     push!(statuses[symbol(node)], st)
 
@@ -209,17 +214,23 @@ julia> b
 2.0
 ```
 """
-function status_from_template(d::Dict{Symbol,T} where {T})
+function status_from_template(d::Dict{Symbol,T}  where {T}, scale::String, orchestrator::Orchestrator2)
     # Sort vars to put the values that are of type PerStatusRef at the end (we need the pass on the other ones first):
     sorted_vars = Dict{Symbol,Any}(sort([pairs(d)...], by=v -> last(v) isa RefVariable ? 1 : 0))
     # Note: PerStatusRef are used to reference variables in the same status for renaming.
 
     # We create the status with the right references for variables, and for PerStatusRef (we reference the reference variable):
     for (k, v) in sorted_vars
-        if isa(v, RefVariable)
-            sorted_vars[k] = sorted_vars[v.reference_variable]
-        else
+        if is_timestep_mapped((scale => k), orchestrator, search_inputs_only=true)
+            # avoid referring to the original variable
             sorted_vars[k] = ref_var(v)
+        else
+
+            if isa(v, RefVariable)
+                sorted_vars[k] = sorted_vars[v.reference_variable]
+            else
+                sorted_vars[k] = ref_var(v)
+            end
         end
     end
 
@@ -423,17 +434,22 @@ function filter_timestep_mapped_variables!(vars_need_init, orchestrator)
 end=#
 
 function filter_timestep_mapped_variables!(vars_need_init, orchestrator)
-    for (org, vars) in vars_need_init
-        for tmst in orchestrator.non_default_timestep_mapping
+    for tmst in orchestrator.non_default_timestep_mapping
+        for (org, vars) in vars_need_init
             if tmst.scale == org
-            
-                for (var_to, var_from) in tmst.var_to_var            
-                    filter!(o -> o == var_to.name || o == var_from.name, vars)
-                end
-                if isempty(vars)
-                    delete!(vars_need_init, org)
+                for (var_from, var_to) in tmst.var_to_var
+                    filter!(o -> o == var_to.name, vars)
+                end                
+            end
+            for (var_from, var_to) in tmst.var_to_var
+                if var_from.scale == org
+                    filter!(o -> o == var_from.name, vars)
                 end
             end
+
+            if isempty(vars)
+                delete!(vars_need_init, org)
+            end
         end
     end
 end

src/mtg/mapping/compute_mapping.jl

@@ -279,14 +279,26 @@ function default_variables_from_mapping(mapped_vars, verbose=true)
 end
 
 
+function is_timestep_mapped(key, orchestrator::Orchestrator2; search_inputs_only::Bool=false)
+    for mtsm in orchestrator.non_default_timestep_mapping
+        for (var_from, var_to) in mtsm.var_to_var
+            if ((first(key) == mtsm.scale) && last(key) == var_to.name) ||
+                (!search_inputs_only && (first(key) == var_from.scale && last(key) == var_from.name))
+                return true
+            end
+        end
+    end
+    return false
+end
+
 """
     convert_reference_values!(mapped_vars::Dict{String,Dict{Symbol,Any}})
 
 Convert the variables that are `MappedVar{SelfNodeMapping}` or `MappedVar{SingleNodeMapping}` to RefValues that reference a 
 common value for the variable; and convert `MappedVar{MultiNodeMapping}` to RefVectors that reference the values for the
 variable in the source organs.
 """
-function convert_reference_values!(mapped_vars::Dict{String,Dict{Symbol,Any}})
+function convert_reference_values!(mapped_vars::Dict{String,Dict{Symbol,Any}}, orchestrator::Orchestrator2)
     # For the variables that will be RefValues, i.e. referencing a value that exists for different scales, we need to first 
     # create a common reference to the value that we use wherever we need this value. These values are created in the dict_mapped_vars
     # Dict, and then referenced from there every time we point to it.
@@ -303,7 +315,11 @@ function convert_reference_values!(mapped_vars::Dict{String,Dict{Symbol,Any}})
 
                 # First time we encounter this variable as a MappedVar, we create its value into the dict_mapped_vars Dict:
                 if !haskey(dict_mapped_vars, key)
-                    push!(dict_mapped_vars, key => Ref(mapped_default(vars[k])))
+#                    if is_timestep_mapped(key, orchestrator)
+#                        push!(dict_mapped_vars, key => (mapped_default(vars[k])))
+#                    else
+                        push!(dict_mapped_vars, key => Ref(mapped_default(vars[k])))
+#                    end
                 end
 
                 # Then we use the value for the particular variable to replace the MappedVar to a RefValue in the mapping:

src/mtg/save_results.jl

@@ -109,8 +109,8 @@ julia> collect(keys(preallocated_vars["Leaf"]))
  :carbon_demand
 ```
 """
-
-function pre_allocate_outputs(statuses, statuses_template, reverse_multiscale_mapping, vars_need_init, outs, nsteps; type_promotion=nothing, check=true)
+# TODO orchestrator prob shouldn't be a kwarg with a default
+function pre_allocate_outputs(statuses, statuses_template, reverse_multiscale_mapping, vars_need_init, outs, nsteps; type_promotion=nothing, check=true, orchestrator=Orchestrator2())
     outs_ = Dict{String,Vector{Symbol}}()
 
     # default behaviour : track everything
@@ -211,18 +211,18 @@ function pre_allocate_outputs(statuses, statuses_template, reverse_multiscale_ma
 
     # I don't know if this function barrier is necessary
     preallocated_outputs = Dict{String,Vector}()
-    complete_preallocation_from_types!(preallocated_outputs, nsteps, outs_, node_type, statuses_template)
+    complete_preallocation_from_types!(preallocated_outputs, nsteps, outs_, node_type, statuses_template, orchestrator)
     return preallocated_outputs
 end
 
-function complete_preallocation_from_types!(preallocated_outputs, nsteps, outs_, node_type, statuses_template)
+function complete_preallocation_from_types!(preallocated_outputs, nsteps, outs_, node_type, statuses_template, orchestrator)
     types = Vector{DataType}()
     for organ in keys(outs_)
 
         outs_no_node = filter(x -> x != :node, outs_[organ])
 
         #types = [typeof(status_from_template(statuses_template[organ])[var]) for var in outs[organ]]
-        values = [status_from_template(statuses_template[organ])[var] for var in outs_no_node]
+        values = [status_from_template(statuses_template[organ], organ, orchestrator)[var] for var in outs_no_node]
 
         #push!(types, node_type)