Merge #1212 #1268

1212: Try using Arrays in RF model r=charleskawczynski a=charleskawczynski This is an alternative to #1211. This PR uses `Array`s in the Random Feature model, since large SArrays and Tuples (O(1000)) are _really_ expensive for the compiler. This PR tests out ordinary arrays. It'll be interesting to compare both the compile time and runtime. 1268: Upgrade thermodynamics r=charleskawczynski a=charleskawczynski This PR upgrades to the new version of thermodynamics, which has some performance improvements. Peel off from #1261. Co-authored-by: Charles Kawczynski <kawczynski.charles@gmail.com>
CliMA · Aug 21, 2022 · 81b0a41 · 81b0a41
3 parents 3849763 + c82efd7 + 43cbe1c
commit 81b0a41
Show file tree

Hide file tree

Showing 7 changed files with 154 additions and 144 deletions.
diff --git a/docs/Manifest.toml b/docs/Manifest.toml
@@ -1631,9 +1631,9 @@ version = "1.0.1"
 
 [[deps.Thermodynamics]]
 deps = ["DocStringExtensions", "KernelAbstractions", "Random", "RootSolvers"]
-git-tree-sha1 = "0ff7428af31cc2925d29384144d495aef19b9047"
+git-tree-sha1 = "0151b1a971fdd9438a6341ffde927efecfe275c1"
 uuid = "b60c26fb-14c3-4610-9d3e-2d17fe7ff00c"
-version = "0.9.3"
+version = "0.9.5"
 
 [[deps.ThreadingUtilities]]
 deps = ["ManualMemory"]

diff --git a/driver/main.jl b/driver/main.jl
@@ -145,9 +145,11 @@ function Simulation1d(namelist)
     end
 
     edmf = TC.EDMFModel(FTD, namelist, precip_model, rain_formation_model)
-    if isbits(edmf)
-        @info "edmf = \n$(summary(edmf))"
-    else
+    # RF contains a very large number of parameters,
+    # using SVectors / Tuples for this is very expensive
+    # for the compiler, so we'll just accept Arrays for now.
+    @info "edmf = \n$(summary(edmf))"
+    if !isbits(edmf) && !(edmf.entr_closure isa TC.RFEntr)
         @show edmf
         error("Something non-isbits was added to edmf and needs to be fixed.")
     end

diff --git a/integration_tests/Manifest.toml b/integration_tests/Manifest.toml
@@ -1700,9 +1700,9 @@ version = "1.0.1"
 
 [[deps.Thermodynamics]]
 deps = ["DocStringExtensions", "KernelAbstractions", "Random", "RootSolvers"]
-git-tree-sha1 = "0ff7428af31cc2925d29384144d495aef19b9047"
+git-tree-sha1 = "0151b1a971fdd9438a6341ffde927efecfe275c1"
 uuid = "b60c26fb-14c3-4610-9d3e-2d17fe7ff00c"
-version = "0.9.3"
+version = "0.9.5"
 
 [[deps.ThreadingUtilities]]
 deps = ["ManualMemory"]

diff --git a/perf/Manifest.toml b/perf/Manifest.toml
@@ -1718,9 +1718,9 @@ version = "1.0.1"
 
 [[deps.Thermodynamics]]
 deps = ["DocStringExtensions", "KernelAbstractions", "Random", "RootSolvers"]
-git-tree-sha1 = "0ff7428af31cc2925d29384144d495aef19b9047"
+git-tree-sha1 = "0151b1a971fdd9438a6341ffde927efecfe275c1"
 uuid = "b60c26fb-14c3-4610-9d3e-2d17fe7ff00c"
-version = "0.9.3"
+version = "0.9.5"
 
 [[deps.ThreadingUtilities]]
 deps = ["ManualMemory"]

diff --git a/post_processing/mse_tables.jl b/post_processing/mse_tables.jl
diff --git a/src/closures/nondimensional_exchange_functions.jl b/src/closures/nondimensional_exchange_functions.jl
@@ -299,16 +299,19 @@ function non_dimensional_function(εδ_model::RFEntr{d, m}, εδ_model_vars) whe
     # Learnable and fixed parameters
     c_rf_fix = εδ_model.c_rf_fix # 2 x m x (1 + d), fix
     c_rf_opt = εδ_model.c_rf_opt # 2 x (m + 1 + d), learn
+    FT = eltype(εδ_model)
+    sqrt2 = FT(sqrt(2))
+    sqrtm = FT(sqrt(m))
 
     # Random Features
     scale_x_entr = (c_rf_opt[1, (m + 2):(m + d + 1)] .^ 2) .* nondim_groups
     scale_x_detr = (c_rf_opt[2, (m + 2):(m + d + 1)] .^ 2) .* nondim_groups
-    f_entr = c_rf_opt[1, m + 1]^2 * sqrt(2) * cos.(c_rf_fix[1, :, 2:(d + 1)] * scale_x_entr + c_rf_fix[1, :, 1])
-    f_detr = c_rf_opt[2, m + 1]^2 * sqrt(2) * cos.(c_rf_fix[2, :, 2:(d + 1)] * scale_x_detr + c_rf_fix[2, :, 1])
+    f_entr = c_rf_opt[1, m + 1]^2 * sqrt2 * cos.(c_rf_fix[1, :, 2:(d + 1)] * scale_x_entr + c_rf_fix[1, :, 1])
+    f_detr = c_rf_opt[2, m + 1]^2 * sqrt2 * cos.(c_rf_fix[2, :, 2:(d + 1)] * scale_x_detr + c_rf_fix[2, :, 1])
 
     # Square output for nonnegativity for prediction
-    nondim_ε = sum(c_rf_opt[1, 1:m] .* f_entr) / sqrt(m)
-    nondim_δ = sum(c_rf_opt[2, 1:m] .* f_detr) / sqrt(m)
+    nondim_ε = sum(c_rf_opt[1, 1:m] .* f_entr) / sqrtm
+    nondim_δ = sum(c_rf_opt[2, 1:m] .* f_detr) / sqrtm
     return Flux.relu(nondim_ε), Flux.relu(nondim_δ)
 end
 

diff --git a/src/types.jl b/src/types.jl
@@ -86,20 +86,25 @@ Base.@kwdef struct FNOEntr{P, T} <: AbstractNonLocalEntrDetrModel
     nm_fno::Int
     c_fno::T
 end
-struct RFEntr{d, m, P, A, B} <: AbstractEntrDetrModel
+struct RFEntr{d, m, FT, P, A, B} <: AbstractEntrDetrModel
     params::P
     c_rf_fix::A
     c_rf_opt::B
     function RFEntr(params::P, c_rf_fix::A, c_rf_opt::B, d::Int) where {P, A, B}
         c_rf_fix = reshape(c_rf_fix, 2, :, 1 + d) # 2 x m x (1 + d), fix
         c_rf_fix = SA.SArray{Tuple{size(c_rf_fix)...}}(c_rf_fix)
+        FT = eltype(c_rf_fix)
+        c_rf_fix = Array{FT}(c_rf_fix)
         m = size(c_rf_fix, 2)
         c_rf_opt = reshape(c_rf_opt, 2, m + 1 + d) # 2 x (m + 1 + d), learn
         c_rf_opt = SA.SArray{Tuple{size(c_rf_opt)...}}(c_rf_opt)
-        return new{d, m, P, typeof(c_rf_fix), typeof(c_rf_opt)}(params, c_rf_fix, c_rf_opt)
+        c_rf_opt = Array{FT}(c_rf_opt)
+        return new{d, m, FT, P, typeof(c_rf_fix), typeof(c_rf_opt)}(params, c_rf_fix, c_rf_opt)
     end
 end
 
+Base.eltype(::RFEntr{d, m, FT}) where {d, m, FT} = FT
+
 Base.@kwdef struct NoisyRelaxationProcess{MT, T} <: AbstractNoisyEntrDetrModel
     mean_model::MT
     c_gen_stoch::T