Merge pull request #39 from PALEOtoolkit/simplify_modeldata

Sync with PALEOboxes v0.21.7 ModelData simplification

Project.toml

@@ -1,7 +1,7 @@
 name = "PALEOmodel"
 uuid = "bf7b4fbe-ccb1-42c5-83c2-e6e9378b660c"
 authors = ["Stuart Daines <[email protected]>"]
-version = "0.15.18"
+version = "0.15.19"
 
 [deps]
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
@@ -36,7 +36,7 @@ ForwardDiff = "0.10"
 Infiltrator = "1.0"
 JLD2 = "0.4"
 NLsolve = "4.5"
-PALEOboxes = "0.21.5"
+PALEOboxes = "0.21.7"
 RecipesBase = "1.2"
 Requires = "1.0"
 Revise = "3.1"

docs/src/PALEOmodelSolvers.md

@@ -116,7 +116,7 @@ CurrentModule = PALEOmodel
 A [`SolverView`](@ref) uses a collection of `PALEOboxes.VariableAggregator`s to assemble model state Variables and associated time derivatives into contiguous Vectors, for the convenience of standard numerical ODE / DAE solvers.  See [Mathematical formulation of the reaction-transport problem](@ref). 
 ```@docs
 SolverView
-create_solver_view
+SolverView
 set_default_solver_view!
 copy_norm!
 set_statevar!

src/ODE.jl

@@ -43,14 +43,12 @@ Keyword arguments are required to generate a Jacobian function (using automatic
 - `jac_ad=:NoJacobian`: Jacobian to generate and use (:NoJacobian, :ForwardDiffSparse, :ForwardDiff)
 - `initial_state::AbstractVector`: initial state vector
 - `jac_ad_t_sparsity::Float64`: model time at which to calculate Jacobian sparsity pattern
-- `init_logger=Logging.NullLogger()`: default value omits logging from (re)initialization to generate Jacobian modeldata, Logging.CurrentLogger() to include
 """
 function ODEfunction(
     model::PB.Model, modeldata;   
     jac_ad=:NoJacobian,
     initial_state=nothing,
     jac_ad_t_sparsity=nothing,    
-    init_logger=Logging.NullLogger(),
     generated_dispatch=true,
 )
     PB.check_modeldata(model, modeldata)
@@ -74,7 +72,7 @@ function ODEfunction(
 
     jac, jac_prototype = PALEOmodel.JacobianAD.jac_config_ode(
         jac_ad, model, initial_state, modeldata, jac_ad_t_sparsity;
-        init_logger, generated_dispatch,
+        generated_dispatch,
     )    
 
     f = SciMLBase.ODEFunction{true}(m; jac, jac_prototype, mass_matrix=M)
@@ -94,14 +92,12 @@ Keyword arguments are required to generate a Jacobian function (using automatic
 - `jac_ad=:NoJacobian`: Jacobian to generate and use (:NoJacobian, :ForwardDiffSparse, :ForwardDiff)
 - `initial_state::AbstractVector`: initial state vector
 - `jac_ad_t_sparsity::Float64`: model time at which to calculate Jacobian sparsity pattern
-- `init_logger=Logging.NullLogger()`: default value omits logging from (re)initialization to generate Jacobian modeldata, Logging.CurrentLogger() to include
 """
 function DAEfunction(
     model::PB.Model, modeldata;   
     jac_ad=:NoJacobian,
     initial_state=nothing,
     jac_ad_t_sparsity=nothing,    
-    init_logger=Logging.NullLogger(),
     generated_dispatch=true,
 )
     @info "DAEfunction:  using Jacobian $jac_ad"
@@ -110,7 +106,7 @@ function DAEfunction(
 
     jac, jac_prototype, odeimplicit = PALEOmodel.JacobianAD.jac_config_dae(
         jac_ad, model, initial_state, modeldata, jac_ad_t_sparsity;
-        init_logger, generated_dispatch,
+        generated_dispatch,
     )
     m =  SolverFunctions.ModelDAE(modeldata, modeldata.solver_view_all, modeldata.dispatchlists_all, odeimplicit, 0)
 
@@ -161,7 +157,6 @@ and then
 - `steadystate=false`: true to use `DifferentialEquations.jl` `SteadyStateProblem`
   (not recommended, see [`PALEOmodel.SteadyState.steadystate`](@ref)).
 - `BLAS_num_threads=1`: number of LinearAlgebra.BLAS threads to use
-- `init_logger=Logging.NullLogger()`: default value omits logging from (re)initialization to generate Jacobian modeldata, Logging.CurrentLogger() to include
 - `generated_dispatch=true`: `true` to autogenerate code (fast solve, slow compile)
 """
 function integrate(
@@ -173,7 +168,6 @@ function integrate(
     jac_ad_t_sparsity=tspan[1],
     steadystate=false,
     BLAS_num_threads=1,
-    init_logger=Logging.NullLogger(),
     generated_dispatch=true,
 )
 
@@ -182,7 +176,6 @@ function integrate(
         jac_ad,
         initial_state,
         jac_ad_t_sparsity,    
-        init_logger,
         generated_dispatch,
     )
 
@@ -269,7 +262,6 @@ function integrateDAE(
     jac_ad=:NoJacobian,
     jac_ad_t_sparsity=tspan[1],
     BLAS_num_threads=1,
-    init_logger=Logging.NullLogger(),
     generated_dispatch=true,
 )
 
@@ -278,7 +270,6 @@ function integrateDAE(
         jac_ad,
         initial_state,
         jac_ad_t_sparsity,    
-        init_logger,
         generated_dispatch,
     )
 

src/ODELocalIMEX.jl

@@ -62,15 +62,15 @@ function integrateLocalIMEXEuler(
 
     PB.check_modeldata(run.model, modeldata)
 
-    solver_view_outer = PALEOmodel.create_solver_view(run.model, modeldata, cellranges_outer)
+    solver_view_outer = PALEOmodel.SolverView(run.model, modeldata, 1, cellranges_outer)
     @info "solver_view_outer: $(solver_view_outer)"    
 
     lictxt = create_timestep_LocalImplicit_ctxt(
         run.model, modeldata;                                   
         cellrange=cellrange_inner,
-        exclude_var_nameroots=exclude_var_nameroots,
-        niter_max=niter_max,
-        Lnorm_inf_max=Lnorm_inf_max
+        exclude_var_nameroots,
+        niter_max,
+        Lnorm_inf_max,
     )
 
     timesteppers = [
@@ -185,20 +185,22 @@ function create_timestep_LocalImplicit_ctxt(
     Lnorm_inf_max,
     request_adchunksize=ForwardDiff.DEFAULT_CHUNK_THRESHOLD,
     generated_dispatch=true,
-    init_logger=Logging.NullLogger(),
 )
     PB.check_modeldata(model, modeldata)
 
     # create SolverViews for first cell, to work out how many dof we need
     cellrange_cell = PB.CellRange(cellrange.domain, cellrange.operatorID, first(cellrange.indices) )
 
     # get SolverView for all Variables required for cellrange derivative
-    solver_view_all_cell = PALEOmodel.create_solver_view(model, modeldata, [cellrange_cell], indices_from_cellranges=true)
+    solver_view_all_cell = PALEOmodel.SolverView(
+        model, modeldata, 1, [cellrange_cell];
+        indices_from_cellranges=true
+    )
     @info "getLocalImplicitContext: all Variables (first cell) $solver_view_all_cell)"  
 
     # get reduced set of Variables required for nonlinear solve
-    solver_view_cell = PALEOmodel.create_solver_view(
-        model, modeldata, [cellrange_cell], 
+    solver_view_cell = PALEOmodel.SolverView(
+        model, modeldata, 1, [cellrange_cell];
         exclude_var_nameroots=exclude_var_nameroots,
         indices_from_cellranges=true,
     )
@@ -210,11 +212,11 @@ function create_timestep_LocalImplicit_ctxt(
     excluded_vars = [
         v for v in solver_view_all_cell.stateexplicit.vars if !(v in solver_view_cell.stateexplicit.vars)
     ]
-    va_excluded = PB.VariableAggregator(excluded_vars, [cellrange for v in excluded_vars], modeldata)
+    va_excluded = PB.VariableAggregator(excluded_vars, [cellrange for v in excluded_vars], modeldata, 1)
     excluded_sms_vars = [
         v for v in solver_view_all_cell.stateexplicit_deriv.vars if !(v in solver_view_cell.stateexplicit_deriv.vars)
     ]
-    va_sms_excluded = PB.VariableAggregator(excluded_sms_vars, [cellrange for v in excluded_sms_vars], modeldata)
+    va_sms_excluded = PB.VariableAggregator(excluded_sms_vars, [cellrange for v in excluded_sms_vars], modeldata, 1)
     length(excluded_vars) == length(excluded_sms_vars) ||
         error("excluded_vars and excluded_sms_vars length mismatch")
 
@@ -226,13 +228,13 @@ function create_timestep_LocalImplicit_ctxt(
     @info "  solve vars: $([PB.fullname(v) for v in solver_view_cell.stateexplicit.vars]) sms_vars: $([PB.fullname(v) for v in solver_view_cell.stateexplicit_deriv.vars])"
     @info "  excluded vars: $([PB.fullname(v) for v in excluded_vars]) sms_vars: $([PB.fullname(v) for v in excluded_sms_vars])"
 
-    # create a modeldata_ad with Dual numbers for AD Jacobians
+    # Add an array set with Dual numbers to modeldata, for AD Jacobians
     chunksize = ForwardDiff.pickchunksize(n_solve, request_adchunksize)
     @info "  using ForwardDiff dense Jacobian chunksize=$chunksize"
-
-    _, modeldata_ad = Logging.with_logger(init_logger) do
-        PALEOmodel.initialize!(model, eltype=ForwardDiff.Dual{Nothing, eltype(modeldata), chunksize}; create_dispatchlists_all=false)
-    end
+    eltype_base = eltype(modeldata, 1)
+    eltype_jac_cell = ForwardDiff.Dual{Nothing, eltype_base, chunksize}
+    PB.add_arrays_data!(model, modeldata, eltype_jac_cell, "jac_cell")
+    arrays_idx_jac_cell = PB.num_arrays(modeldata)
 
     # temporarily replace modeldata with norm so can read back per-cell norms
     statevar_all_current = PALEOmodel.get_statevar(modeldata.solver_view_all)
@@ -248,9 +250,9 @@ function create_timestep_LocalImplicit_ctxt(
         # (ab)use that fact that Julia allows iteration over scalar i (to optimise out loop over cellrange.indices)
         cellrange = PB.CellRange(cellrange.domain, cellrange.operatorID, i) 
 
-        sv = PALEOmodel.create_solver_view(
-            model, modeldata, [cellrange], 
-            exclude_var_nameroots=exclude_var_nameroots,
+        sv = PALEOmodel.SolverView(
+            model, modeldata, 1, [cellrange];
+            exclude_var_nameroots,
             indices_from_cellranges=true,
             hostdep_all=false,
         )
@@ -260,18 +262,18 @@ function create_timestep_LocalImplicit_ctxt(
         statevar_norm = PALEOmodel.get_statevar_norm(sv)
         push!(statevar_norms, statevar_norm)        
 
-        dl = PB.create_dispatch_methodlists(model, modeldata, [cellrange]; generated_dispatch)
+        dl = PB.create_dispatch_methodlists(model, modeldata, 1, [cellrange]; generated_dispatch)
         push!(dispatchlists, dl)
 
-        sv_ad = PALEOmodel.create_solver_view(
-            model, modeldata_ad, [cellrange], 
-            exclude_var_nameroots=exclude_var_nameroots,
+        sv_ad = PALEOmodel.SolverView(
+            model, modeldata, arrays_idx_jac_cell, [cellrange]; 
+            exclude_var_nameroots,
             indices_from_cellranges=true,
             hostdep_all=false,
         )
         push!(solverviews_ad, sv_ad)
 
-        dl_ad = PB.create_dispatch_methodlists(model, modeldata_ad, [cellrange]; generated_dispatch)
+        dl_ad = PB.create_dispatch_methodlists(model, modeldata, arrays_idx_jac_cell, [cellrange]; generated_dispatch)
 
         push!(dispatchlists_ad, dl_ad)
 
@@ -282,7 +284,7 @@ function create_timestep_LocalImplicit_ctxt(
 
     cell_residual = CellResidual(
         n_solve,
-        eltype(modeldata),
+        eltype_base,
         cellindices,
         [sv for sv in solverviews],  # narrow type
         [dl for dl in dispatchlists], # narrow type
@@ -292,7 +294,7 @@ function create_timestep_LocalImplicit_ctxt(
     cell_jacobian = CellJacobian(
         CellResidual(
             n_solve,
-            eltype(modeldata_ad),
+            eltype_jac_cell,
             cellindices,
             [sv for sv in solverviews_ad],  # narrow type
             [dl for dl in dispatchlists_ad], # narrow type

src/ODEfixed.jl

@@ -99,9 +99,9 @@ function integrateSplitEuler(
 
     PB.check_modeldata(run.model, modeldata)
 
-    solver_view_outer = PALEOmodel.create_solver_view(run.model, modeldata, cellranges_outer)
+    solver_view_outer = PALEOmodel.SolverView(run.model, modeldata, 1, cellranges_outer)
     @info "solver_view_outer: $(solver_view_outer)"    
-    solver_view_inner = PALEOmodel.create_solver_view(run.model, modeldata, cellranges_inner)
+    solver_view_inner = PALEOmodel.SolverView(run.model, modeldata, 1, cellranges_inner)
     @info "solver_view_inner: $(solver_view_inner)"
 
     timesteppers = [
@@ -174,7 +174,7 @@ function integrateEulerthreads(
         error("integrateEulerthreads: length(cellranges) $lc != nthreads $nt")
 
     # get solver_views for each threadid
-    solver_views = [PALEOmodel.create_solver_view(run.model, modeldata, crs) for crs in cellranges]
+    solver_views = [PALEOmodel.SolverView(run.model, modeldata, 1, crs) for crs in cellranges]
     @info "integrateEulerthreads: solver_views:" 
     for t in 1:Threads.nthreads()
         @info "  thread $t  $(solver_views[t])"
@@ -253,8 +253,8 @@ function integrateSplitEulerthreads(
         error("integrateSplitEulerthreads: length(cellranges_inner) $lc_inner != nthreads $nt")
 
     # get solver_views for each threadid
-    solver_views_outer = [PALEOmodel.create_solver_view(run.model, modeldata, crs) for crs in cellranges_outer]
-    solver_views_inner = [PALEOmodel.create_solver_view(run.model, modeldata, crs) for crs in cellranges_inner]
+    solver_views_outer = [PALEOmodel.SolverView(run.model, modeldata, 1, crs) for crs in cellranges_outer]
+    solver_views_inner = [PALEOmodel.SolverView(run.model, modeldata, 1, crs) for crs in cellranges_inner]
     @info "integrateSplitEulerthreads: solver_views_outer:" 
     for t in 1:Threads.nthreads()
         @info "  thread $t $(solver_views_outer[t])"
@@ -349,7 +349,7 @@ function create_timestep_Euler_ctxt(
     num_constraints = PALEOmodel.num_algebraic_constraints(solver_view)
     iszero(num_constraints) || error("DAE problem with $num_constraints algebraic constraints")
 
-    dispatch_lists = PB.create_dispatch_methodlists(model, modeldata, cellranges; verbose, generated_dispatch)
+    dispatch_lists = PB.create_dispatch_methodlists(model, modeldata, 1, cellranges; verbose, generated_dispatch)
 
     return (dispatch_lists, solver_view, n_substep)
 end

src/Run.jl

@@ -64,7 +64,7 @@ and supplying `method_barrier` (a thread barrier to add to `ReactionMethod` disp
 - `threadsafe=false`: true to create thread safe Atomic Variables where Variable attribute `:atomic==true`
 - `method_barrier=nothing`: thread barrier to add to dispatch lists if `threadsafe==true`
 - `expect_hostdep_varnames=["global.tforce"]`: non-state-Variable host-dependent Variable names expected
-- `create_solver_view_all=true`: `true` to create `modeldata.solver_view_all`
+- `SolverView_all=true`: `true` to create `modeldata.solver_view_all`
 - `create_dispatchlists_all=true`: `true` to create `modeldata.dispatchlists_all`
 - `generated_dispatch=true`: `true` to autogenerate code for `modeldata.dispatchlists_all` (fast dispatch, slow compile)
 """
@@ -81,21 +81,21 @@ function initialize!(
         ) : 
         nothing,
     expect_hostdep_varnames=["global.tforce"],
-    create_solver_view_all=true,
+    SolverView_all=true,
     create_dispatchlists_all=true,
     generated_dispatch=true,
 )
 
     modeldata = PB.create_modeldata(model, eltype; threadsafe)
 
     # Allocate variables
-    @timeit "allocate_variables" PB.allocate_variables!(model, modeldata; eltypemap)
+    @timeit "allocate_variables" PB.allocate_variables!(model, modeldata, 1; eltypemap)
 
     # check all variables allocated
     PB.check_ready(model, modeldata; expect_hostdep_varnames)
 
     # Create modeldata.solver_view_all for the entire model
-    if create_solver_view_all
+    if SolverView_all
         @timeit "set_default_solver_view!" set_default_solver_view!(model, modeldata)
     end
 

src/SolverFunctions.jl

@@ -27,8 +27,8 @@ Function object to calculate model time derivative and adapt to SciML ODE solver
 
 Call as `f(du,u, p, t)`
 """
-mutable struct ModelODE{T, S <: PALEOmodel.SolverView, D}
-    modeldata::PB.ModelData{T}
+mutable struct ModelODE{S <: PALEOmodel.SolverView, D}
+    modeldata::PB.ModelData
     solver_view::S
     dispatchlists::D
     nevals::Int
@@ -224,8 +224,8 @@ where residual `G(dsdt,s,p,t)` is:
 - `F(s)` (for algebraic constraints)
 - `duds*dsdt + F(s, u(s))` (for Total variables u that depend implicitly on state Variables s)
 """
-mutable struct ModelDAE{T, S <: PALEOmodel.SolverView, D, O}
-    modeldata::PB.ModelData{T}
+mutable struct ModelDAE{S <: PALEOmodel.SolverView, D, O}
+    modeldata::PB.ModelData
     solver_view::S
     dispatchlists::D
     odeimplicit::O
@@ -357,8 +357,8 @@ end
 
 Calculate dT/dS required for a model containing implicit Total variables, using ForwardDiff and dense AD
 """
-mutable struct ImplicitForwardDiffDense{T, S, D, W, I, U}
-    modeldata::PB.ModelData{T}
+mutable struct ImplicitForwardDiffDense{S, D, W, I, U}
+    modeldata::PB.ModelData
     implicitderiv::TotalForwardDiff{S, D}
     duds_template::W
     implicitconf::I
@@ -388,8 +388,8 @@ end
 Calculate dT/dS required for a model containing implicit Total variables, using ForwardDiff and 
 sparse AD with `SparseDiffTools.forwarddiff_color_jacobian!`
 """
-mutable struct ImplicitForwardDiffSparse{T, S, D, I, U}
-    modeldata::PB.ModelData{T}
+mutable struct ImplicitForwardDiffSparse{S, D, I, U}
+    modeldata::PB.ModelData
     implicitderiv::TotalForwardDiff{S, D}
     implicit_cache::I
     duds::U