Better defaults for WENOVectorInvariant (#3365)

* Add better default for WENOVectorInvariant * Find max order among all schemes
CliMA · Oct 30, 2023 · f3996c2 · f3996c2 · navidcy · Nov 1, 2023
1 parent 19cead9
commit f3996c2
Showing 1 changed file with 28 additions and 10 deletions.
diff --git a/src/Advection/vector_invariant_advection.jl b/src/Advection/vector_invariant_advection.jl
@@ -170,6 +170,8 @@ Base.show(io::IO, a::VectorInvariant{N, FT}) where {N, FT} =
 #                           VectorInvariant{N,     FT,    M,     Z (vorticity scheme),      ZS,    V (vertical scheme),      K (kinetic energy gradient scheme)
 const WENOVectorInvariant = VectorInvariant{<:Any, <:Any, <:Any, <:WENO, <:Any, <:WENO, <:WENO}
 
+nothing_to_default(user_value, default) = isnothing(user_value) ? default : user_value
+
 """
     function WENOVectorInvariant(; upwinding = nothing,
                                    multi_dimensional_stencil = false,
@@ -178,21 +180,37 @@ const WENOVectorInvariant = VectorInvariant{<:Any, <:Any, <:Any, <:WENO, <:Any,
 """
 function WENOVectorInvariant(; upwinding = nothing,
                                vorticity_stencil = VelocityStencil(),
+                               order = nothing,
+                               vorticity_order = nothing,
+                               vertical_order = nothing,
+                               divergence_order = nothing,
+                               kinetic_energy_gradient_order = nothing, 
                                multi_dimensional_stencil = false,
                                weno_kw...)
 
-    weno = WENO(; weno_kw...)
-    vorticity_scheme = weno
-    vertical_scheme = weno
-    kinetic_energy_gradient_scheme = weno
-    divergence_scheme = weno
-
-    if isnothing(upwinding)
-        upwinding  = OnlySelfUpwinding(; cross_scheme = weno)
+    if isnothing(order) # apply global defaults
+        vorticity_order               = nothing_to_default(vorticity_order, default=9)
+        vertical_order                = nothing_to_default(vertical_order, default=5)
+        divergence_order              = nothing_to_default(divergence_order, default=5)
+        kinetic_energy_gradient_order = nothing_to_default(kinetic_energy_gradient_order, default=5)
+    else # apply user supplied `order` unless overridden by more specific value
+        vorticity_order               = nothing_to_default(vorticity_order, default=order)
+        vertical_order                = nothing_to_default(vertical_order, default=order)
+        divergence_order              = nothing_to_default(divergence_order, default=order)
+        kinetic_energy_gradient_order = nothing_to_default(kinetic_energy_gradient_order, default=order)
     end
 
-    N = required_halo_size(weno)
-    FT = eltype(weno)
+    vorticity_scheme               = WENO(; order=vorticity_order, weno_kw...)
+    vertical_scheme                = WENO(; order=vertical_order, weno_kw...)
+    kinetic_energy_gradient_scheme = WENO(; order=kinetic_energy_gradient_order, weno_kw...)
+    divergence_scheme              = WENO(; order=divergence_order, weno_kw...)
+
+    default_upwinding = OnlySelfUpwinding(cross_scheme=divergence_scheme)
+    upwinding = nothing_to_default(upwinding; default = default_upwinding)
+
+    schemes = (vorticity_scheme, vertical_scheme, kinetic_energy_gradient_scheme, divergence_scheme)
+    N = maximum(required_halo_size(s) for s in schemes)
+    FT = eltype(vorticity_scheme) # assumption
 
     return VectorInvariant{N, FT, multi_dimensional_stencil}(vorticity_scheme,
                                                              vorticity_stencil,