refactor: modularize process_SciMLProblem a bit more

Author: AayushSabharwal

src/systems/problem_utils.jl

@@ -1192,6 +1192,20 @@ function float_type_from_varmap(varmap, floatT = Bool)
     return float(floatT)
 end
 
+"""
+    $(TYPEDSIGNATURES)
+
+Calculate the floating point type to use from the given `varmap` by looking at variables
+with a constant value. `u0Type` takes priority if it is a real-valued array type.
+"""
+function calculate_float_type(varmap, u0Type::Type, floatT = Bool)
+    if u0Type <: AbstractArray && eltype(u0Type) <: Real && eltype(u0Type) != Union{}
+        return float(eltype(u0Type))
+    else
+        return float_type_from_varmap(varmap, floatT)
+    end
+end
+
 """
     $(TYPEDSIGNATURES)
 
@@ -1208,6 +1222,41 @@ function calculate_resid_prototype(N::Int, u0, p)
     return zeros(u0ElType, N)
 end
 
+"""
+    $(TYPEDSIGNATURES)
+
+Given the user-provided value of `u0_constructor`, the container type of user-provided
+`op`, the desired floating point type and whether a symbolic `u0` is allowed, return the
+updated `u0_constructor`.
+"""
+function get_u0_constructor(u0_constructor, u0Type::Type, floatT::Type, symbolic_u0::Bool)
+    u0_constructor === identity || return u0_constructor
+    u0Type <: StaticArray || return u0_constructor
+    return function (vals)
+        elT = if symbolic_u0 && any(x -> symbolic_type(x) != NotSymbolic(), vals)
+            nothing
+        else
+            floatT
+        end
+        SymbolicUtils.Code.create_array(u0Type, elT, Val(1), Val(length(vals)), vals...)
+    end
+end
+
+"""
+    $(TYPEDSIGNATURES)
+
+Given the user-provided value of `p_constructor`, the container type of user-provided `op`,
+ans the desired floating point type, return the updated `p_constructor`.
+"""
+function get_p_constructor(p_constructor, pType::Type, floatT::Type)
+    p_constructor === identity || return p_constructor
+    pType <: StaticArray || return p_constructor
+    return function (vals)
+        SymbolicUtils.Code.create_array(
+            pType, floatT, Val(ndims(vals)), Val(size(vals)), vals...)
+    end
+end
+
 """
     $(TYPEDSIGNATURES)
 
@@ -1274,26 +1323,15 @@ function process_SciMLProblem(
     missing_unknowns, missing_pars = build_operating_point!(sys, op,
         u0map, pmap, defs, dvs, ps)
 
-    floatT = Bool
-    if u0Type <: AbstractArray && eltype(u0Type) <: Real && eltype(u0Type) != Union{}
-        floatT = float(eltype(u0Type))
-    else
-        floatT = float_type_from_varmap(op, floatT)
-    end
-
+    floatT = calculate_float_type(op, u0Type)
     u0_eltype = something(u0_eltype, floatT)
 
     if !is_time_dependent(sys) || is_initializesystem(sys)
         add_observed_equations!(op, obs)
     end
-    if u0_constructor === identity && u0Type <: StaticArray
-        u0_constructor = vals -> SymbolicUtils.Code.create_array(
-            u0Type, floatT, Val(1), Val(length(vals)), vals...)
-    end
-    if p_constructor === identity && pType <: StaticArray
-        p_constructor = vals -> SymbolicUtils.Code.create_array(
-            pType, floatT, Val(1), Val(length(vals)), vals...)
-    end
+
+    u0_constructor = get_u0_constructor(u0_constructor, u0Type, u0_eltype, symbolic_u0)
+    p_constructor = get_p_constructor(p_constructor, pType, floatT)
 
     if build_initializeprob
         kws = maybe_build_initialization_problem(