Merge pull request #63 from PALEOtoolkit/simplify_modeldata

ModelData can now store multiple sets of arrays for different element types

Project.toml

@@ -1,7 +1,7 @@
 name = "PALEOboxes"
 uuid = "804b410e-d900-4b2a-9ecd-f5a06d4c1fd4"
 authors = ["Stuart Daines <[email protected]>"]
-version = "0.21.6"
+version = "0.21.7"
 
 [deps]
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"

docs/src/DomainsVariablesFields.md

@@ -145,6 +145,7 @@ CurrentModule = PALEOboxes
 AbstractSpace
 ScalarSpace
 CellSpace
+ColumnSpace
 ```
 
 ## Data types

docs/src/Solver API.md

@@ -15,6 +15,8 @@ CurrentModule = PALEOboxes
 create_model_from_config(config_file::AbstractString, configmodel::AbstractString)
 create_modeldata(model::Model, vectype::DataType=Array{Float64,1})
 ModelData
+add_arrays_data!
+push_arrays_data!
 allocate_variables!
 
 check_ready(model::Model, modeldata::AbstractModelData; throw_on_error=true)

src/Domain.jl

@@ -177,55 +177,52 @@ function get_reaction(domain::Domain, reactname::AbstractString; allow_not_found
 end
 
 """
-    allocate_variables!(domain, modeldata; [hostdep=false] [, eltypemap::Dict{String, DataType}])
+    allocate_variables!(domain, modeldata, arrays_idx; [hostdep=false] [, kwargs...])
 
 Allocate memory for Domain Variables. 
 
 If `hostdep=false`, only internal Variables are allocated, allowing host-dependent Variables
 (usually state Variables and derivatives + any external dependencies) to be set to views on host-managed arrays.
 
-Element type of allocated Arrays is determined by `eltype(modeldata)` (the usual case, allowing use of AD types), 
-or Variable `:datatype` attribute if present (allowing Variables to ignore AD types).
-`:datatype` may be either a Julia `DataType` (eg Float64), or a string to be looked up in `eltypemap`.
+See [`allocate_variables!(vars, modeldata::AbstractModelData, arrays_idx::Int)`](@ref).
 """
 function allocate_variables!(
-    domain::Domain, modeldata::AbstractModelData; 
+    domain::Domain, modeldata::AbstractModelData, arrays_idx::Int; 
     hostdep::Union{Bool,Nothing}=nothing,
-    eltypemap=Dict{String, DataType}()
+    kwargs...
 )
     vars = get_variables(domain, hostdep=hostdep)
 
     @info "Domain $(rpad(domain.name,20)) data dimensions $(rpad(domain.data_dims,20)) "*
         "allocating $(rpad(length(vars),4)) variables (hostdep=$(hostdep))"  
 
     allocate_variables!(
-        vars, modeldata; 
-        eltypemap=eltypemap,
-        default_host_dependent_field_data=ScalarData,
+        vars, modeldata, arrays_idx; 
+        kwargs...
     )
 
     return nothing
 end
 
 """
-    get_unallocated_variables(domain, modeldata) -> Vector{VariableDomain}
+    get_unallocated_variables(domain, modeldata, arrays_idx::Int) -> Vector{VariableDomain}
 
 Return any unallocated variables (host-dependent variables which have no data pointer set)
 """
 function get_unallocated_variables(
-    domain::Domain, modeldata::AbstractModelData
+    domain::Domain, modeldata::AbstractModelData, arrays_idx::Int
 )
     allvars = get_variables(domain)
-    unallocated_variables = [v for v in allvars if !is_allocated(v, modeldata)]
+    unallocated_variables = [v for v in allvars if !is_allocated(v, modeldata, arrays_idx)]
     return unallocated_variables
 end
 
 "Check all variable pointers set"
 function check_ready(
-    domain::Domain, modeldata::AbstractModelData;
+    domain::Domain, modeldata::AbstractModelData, arrays_idx::Int=1;
     throw_on_error=true
 )
-    vars_unallocated = get_unallocated_variables(domain, modeldata)
+    vars_unallocated = get_unallocated_variables(domain, modeldata, arrays_idx)
     num_unallocated = length(vars_unallocated)
     if num_unallocated == 0
         return true

src/Model.jl

@@ -317,15 +317,11 @@ end
 Create a new [`ModelData`](@ref) struct for model variables of element type `eltype`.
 """
 function create_modeldata(
-    model::Model, eltype::DataType=Float64;
+    model::Model, arrays_eltype::DataType=Float64;
     threadsafe=false,
     allocatenans=true, # fill Arrays with NaN when first allocated
 )
-    modeldata = ModelData{eltype}(model=model, threadsafe=threadsafe, allocatenans=allocatenans)
-
-    for dom in model.domains
-        push!(modeldata.domain_data, DomainData(dom))      
-    end
+    modeldata = ModelData(model; arrays_eltype, threadsafe, allocatenans)
 
     modeldata.cellranges_all = create_default_cellrange(model)
 
@@ -348,24 +344,23 @@ function create_default_cellrange(model::Model ; operatorID=0)
 end
 
 """
-    allocate_variables!(model, modeldata; [hostdep] [,eltypemap])
+    allocate_variables!(model, modeldata, arrays_idx; kwargs...)
     
 Allocate memory for Domain variables for every Domain in `model`.
 
-See [`allocate_variables!(domain::Domain, modeldata::AbstractModelData)`](@ref).
+See [`allocate_variables!(domain::Domain, modeldata::AbstractModelData, arrays_idx::Int)`](@ref).
 """
 function allocate_variables!(
-    model::Model, modeldata::AbstractModelData; 
-    hostdep::Union{Bool,Nothing}=nothing,
-    eltypemap=Dict{String, DataType}()
+    model::Model, modeldata::AbstractModelData, arrays_idx::Int; 
+    kwargs...
 )
     @info lpad("", 80, "=")
-    @info "allocate_variables!"
+    @info "allocate_variables! (modeldata arrays_idx=$arrays_idx)"
     @info lpad("", 80, "=")
 
     check_modeldata(model, modeldata)
     for dom in model.domains
-        allocate_variables!(dom, modeldata, hostdep=hostdep, eltypemap=eltypemap)
+        allocate_variables!(dom, modeldata, arrays_idx; kwargs...)
     end
 
     return nothing
@@ -444,17 +439,19 @@ Optionally calls `create_dispatch_methodlists(model, modeldata, modeldata.cellra
 to set `modeldata.dispatchlists_all` to default  ReactionMethodDispatchLists for entire model.
 
 # Keyword arguments
-- `create_dispatchlists_all=true`: true to set `modeldata.dispatchlists_all`
+- `arrays_indices=1:num_arrays(modeldata)`: `modeldata` `arrays_idx` to generate dispatch lists for
+- `create_dispatchlists_all=false`: true to set `modeldata.dispatchlists_all`
 - `generated_dispatch=true`: true to use autogenerated code for `modeldata.dispatchlists_all` (fast dispatch, slow compile)
 """
 function initialize_reactiondata!(
     model::Model, modeldata::AbstractModelData;
+    arrays_indices=1:num_arrays(modeldata),
     method_barrier=nothing,
-    create_dispatchlists_all=true,
+    create_dispatchlists_all=false,
     generated_dispatch=true,
 )
     @info lpad("", 80, "=")
-    @info "initialize_reactiondata!"
+    @info "initialize_reactiondata! (modeldata arrays_indices=$arrays_indices)"
     @info lpad("", 80, "=")
 
     check_modeldata(model, modeldata)
@@ -464,28 +461,38 @@ function initialize_reactiondata!(
     # (Ref gives fast ReactionMethodDispatchList creation, but slow first do_deriv)
     # NB: passing Ref to call_method seems to speed up first do_deriv 
 
-    modeldata.sorted_methodsdata_setup = 
-        [
-            Any[Ref(m), Ref(m.preparefn(m, create_accessors(m, modeldata)))]
-            for m in get_methods(model.sorted_methods_setup)
-        ]
+    # only for arrays_idx = 1
+
 
-    modeldata.sorted_methodsdata_initialize = 
-        [
-            Any[Ref(m), Ref(m.preparefn(m, create_accessors(m, modeldata)))]
-            for m in get_methods(model.sorted_methods_initialize; method_barrier)
-        ]
+    for arrays_idx in arrays_indices
+        arrays_idx in 1:num_arrays(modeldata) || error("arrays_idx $arrays_idx out of range")
+
+        if arrays_idx == 1
+            modeldata.sorted_methodsdata_setup = 
+                [
+                    Any[Ref(m), Ref(m.preparefn(m, create_accessors(m, modeldata, 1)))]
+                    for m in get_methods(model.sorted_methods_setup)
+                ]
+        end
+
+        modeldata.sorted_methodsdata_initialize[arrays_idx] = 
+            [
+                Any[Ref(m), Ref(m.preparefn(m, create_accessors(m, modeldata, arrays_idx)))]
+                for m in get_methods(model.sorted_methods_initialize; method_barrier)
+            ]
+
 
-    modeldata.sorted_methodsdata_do = 
-        [
-            Any[Ref(m), Ref(m.preparefn(m, create_accessors(m, modeldata)))]
-            for m in get_methods(model.sorted_methods_do; method_barrier)
-        ]
+        modeldata.sorted_methodsdata_do[arrays_idx] = 
+            [
+                Any[Ref(m), Ref(m.preparefn(m, create_accessors(m, modeldata, arrays_idx)))]
+                for m in get_methods(model.sorted_methods_do; method_barrier)
+            ]
+    end
 
     if create_dispatchlists_all
-        # create default dispatchlists_all corresponding to cellranges_all
+        # create default dispatchlists_all corresponding to cellranges_all for arrays_idx=1
         modeldata.dispatchlists_all =
-            create_dispatch_methodlists(model, modeldata, modeldata.cellranges_all; generated_dispatch)
+            create_dispatch_methodlists(model, modeldata, 1, modeldata.cellranges_all; generated_dispatch)
     end
 
     return nothing
@@ -498,9 +505,10 @@ end
 Call setup methods, eg to initialize data arrays (including state variables).
 
 `attribute_name` defines the setup operation performed. `dispatch_setup` should be called in sequence with `attribute_name` = :
-- `:setup`: initialise Reactions and set up any non-state Variables (eg model grid Variables)
-- `:norm_value`: set state Variable values from `:norm_value` attribute in .yaml file, and initialise any Reaction state that requires this value.
-- `:initial_value` (optional): set state Variable values from `:initial_value` attribute in .yaml file.
+- `:setup`: initialise Reactions and set up any non-state Variables (eg model grid Variables) (applied to `modeldata` `arrays_idx=1`, values then copied to other `arrays_idx`)
+- `:norm_value`: set state Variable values from `:norm_value` attribute in .yaml file, and initialise any Reaction state that requires this value (`arrays_idx` 1 only)
+- `:initial_value` (optional): set state Variable values from `:initial_value` attribute in .yaml file (`arrays_idx` 1 only)
+
 """
 function dispatch_setup(
     model::Model, attribute_name, modeldata::AbstractModelData, cellranges=modeldata.cellranges_all
@@ -517,11 +525,54 @@ function dispatch_setup(
         end
     end
 
+    if attribute_name == :setup
+        copy_base_values!(modeldata)
+    end
+
     return nothing  
 end
 
+
+
+"""
+    add_arrays_data!(
+        model, modeldata, arrays_eltype::DataType, arrays_tagname::AbstractString;
+        [method_barrier=nothing] [, generated_dispatch=true] [, kwargs...])
+    
+Add a data array set to `modeldata`, allocate memory, and initialize reactiondata.
+
+Element type and tag name are set by `arrays_eltype`, `arrays_tagname`
+
+See [`allocate_variables!(model::Model, modeldata::AbstractModelData, arrays_idx::Int)`](@ref) and 
+[`initialize_reactiondata!`] for keyword arguments.
+"""
+function add_arrays_data!(
+    model::Model, modeldata::AbstractModelData, arrays_eltype::DataType, arrays_tagname::AbstractString;
+    method_barrier=nothing,
+    kwargs...
+)
+    @info lpad("", 80, "=")
+    @info "add_arrays_eltype! (arrays_eltype=$arrays_eltype, arrays_tagname=$arrays_tagname)"
+    @info lpad("", 80, "=")
+
+    push_arrays_data!(modeldata, arrays_eltype, arrays_tagname)
+
+    allocate_variables!(model, modeldata, num_arrays(modeldata); kwargs...)
+
+    copy_base_values!(modeldata, num_arrays(modeldata))
+
+    initialize_reactiondata!(
+        model, modeldata;
+        arrays_indices=num_arrays(modeldata),
+        method_barrier,
+        create_dispatchlists_all=false,
+    )
+
+    return nothing
+end
+
 """
-    create_dispatch_methodlists(model::Model, modeldata::AbstractModelData, cellranges; kwargs) 
+    create_dispatch_methodlists(model::Model, modeldata::AbstractModelData, arrays_idx::Int, cellranges; kwargs) 
         -> (;list_initialize, list_do)
  
 Compile lists of `initialize` and `do` methods + corresponding `cellrange` for main loop [`do_deriv`](@ref).
@@ -534,21 +585,23 @@ Subset of methods and cellrange to operate on are generated from supplied `cellr
   `false` to create `ReactionMethodDispatchListNoGen` (slow dynamic dispatch, fast compile time)
 """
 function create_dispatch_methodlists(
-    model::Model, modeldata::AbstractModelData, cellranges;
+    model::Model, modeldata::AbstractModelData, arrays_idx::Int, cellranges;
     verbose=false,
     generated_dispatch=true,
 )
     check_modeldata(model, modeldata)
 
     verbose && @info "list_initialize:\n"
+    arrays_idx in 1:length(modeldata.sorted_methodsdata_initialize) || error("list_initialize: arrays_idx $arrays_idx not available")
     @timeit "list_initialize" list_initialize = _create_dispatch_methodlist(
-        modeldata.sorted_methodsdata_initialize,
+        modeldata.sorted_methodsdata_initialize[arrays_idx],
         cellranges,
         generated_dispatch,
     )
     verbose && @info "list_do:\n"
+    arrays_idx in 1:length(modeldata.sorted_methodsdata_do) || error("list_do: arrays_idx $arrays_idx not available")
     @timeit "list_do" list_do = _create_dispatch_methodlist(
-        modeldata.sorted_methodsdata_do,
+        modeldata.sorted_methodsdata_do[arrays_idx],
         cellranges,
         generated_dispatch,
     )
@@ -592,6 +645,8 @@ function _dispatch_cellranges(@nospecialize(method::AbstractReactionMethod), cel
     end
 end
 
+
+
 ####################################
 # Main loop methods
 #####################################