diff --git a/src/initial_conditions/initial_conditions.jl b/src/initial_conditions/initial_conditions.jl index 2b4b31e39a2..8b529220935 100644 --- a/src/initial_conditions/initial_conditions.jl +++ b/src/initial_conditions/initial_conditions.jl @@ -20,12 +20,14 @@ struct ColumnInterpolatableField{F, D} f::F data::D function ColumnInterpolatableField(f::Fields.ColumnField) - zdata = parent(Fields.Fields.coordinate_field(f).z) - data = Dierckx.Spline1D(zdata, parent(f); k = 1) + zdata = vec(parent(Fields.Fields.coordinate_field(f).z)) + fdata = vec(parent(f)) + @assert length(zdata) == length(fdata) + data = Dierckx.Spline1D(zdata, fdata; k = 1) return new{typeof(f), typeof(data)}(f, data) end end -(f::ColumnInterpolatableField)(z) = f.data(z) +(f::ColumnInterpolatableField)(z) = Spaces.undertype(axes(f.f))(f.data(z)) #= Solve for x: @@ -34,11 +36,13 @@ Solve for x: (ᶠx^{k+1}-ᶠx^{k})/ᶜΔz = ᶜf(ᶜx̄,ᶜz) ᶜx̄ = (x^{k+1}+x^{k})/2 (ᶠx^{k+1}-ᶠx^{k})/ᶜΔz = ᶜf((ᶠx^{k+1}+ᶠx^{k})/2,ᶜz) -root equation: (_x-x^{k})/Δz = f((x^{k+1}+x^{k})/2,ᶜz) +root equation: (_x-x^{k})/Δz = f((_x+x^{k})/2,ᶜz) ``` =# import RootSolvers +import ClimaComms +import ClimaCore.Utilities: half function column_indefinite_integral!( f::Function, ᶠintegral::Fields.ColumnField, @@ -49,15 +53,15 @@ function column_indefinite_integral!( face_space = axes(ᶠintegral) first_level = Operators.left_idx(face_space) last_level = Operators.right_idx(face_space) - ᶜΔzfield = Fields.Δz_field(ᶜz) + ᶜΔzfield = Fields.Δz_field(ᶜzfield) @inbounds Fields.level(ᶠintegral, first_level)[] = x₀ x₁ = x₀ for level in (first_level + 1):last_level - ᶜz = Fields.level(ᶜzfield, level - half)[] + ᶜz = Fields.level(ᶜzfield.z, level - half)[] ᶜΔz = Fields.level(ᶜΔzfield, level - half)[] x₀ = x₁ function root_eq(_x) - return (_x - x₀) /ᶜ Δz - f(average(_x, x₀), ᶜz) + return (_x - x₀) / ᶜΔz - f(average(_x, x₀), ᶜz) end sol = RootSolvers.find_zero( root_eq, @@ -67,23 +71,31 @@ function column_indefinite_integral!( x₁ = sol.root f₁ = f(average(x₀, x₁), ᶜz) ᶜintegrand_lev = f₁ - @inbounds Fields.level(ᶠintegral, level)[] = radd( + @inbounds Fields.level(ᶠintegral, level)[] = RecursiveApply.radd( Fields.level(ᶠintegral, level - 1)[], - rmul(ᶜintegrand_lev, ᶜΔz), + RecursiveApply.rmul(ᶜintegrand_lev, ᶜΔz), ) end return nothing end +import ClimaCore.Domains as Domains +import ClimaCore.Meshes as Meshes +import ClimaCore.Geometry as Geometry +import ClimaCore.Operators as Operators +import ClimaCore.Topologies as Topologies +import ClimaCore.Spaces as Spaces +import ClimaCore.Fields as Fields +import ClimaCore.RecursiveApply as RecursiveApply """ - column_indefinite_integral(dp_dz, p_0, (FT(0), z_max)) + column_indefinite_integral(f, x₀, zspan; nelems = 100) The column integral, returned as an interpolate-able field. """ function column_indefinite_integral( f::Function, - x₀, + x₀::FT, zspan::Tuple{FT, FT}; nelems = 100, # sets resolution for integration ) where {FT <: Real} @@ -623,15 +635,24 @@ function hydrostatic_pressure_profile(; ts(p, z, T::FunctionOrSpline, θ::FunctionOrSpline, _) = error("Only one of T and θ can be specified") ts(p, z, T::FunctionOrSpline, ::Nothing, ::Nothing) = - TD.PhaseDry_pT(thermo_params, p, T(z)) + TD.PhaseDry_pT(thermo_params, p, oftype(p, T(z))) ts(p, z, ::Nothing, θ::FunctionOrSpline, ::Nothing) = - TD.PhaseDry_pθ(thermo_params, p, θ(z)) + TD.PhaseDry_pθ(thermo_params, p, oftype(p, θ(z))) ts(p, z, T::FunctionOrSpline, ::Nothing, q_tot::FunctionOrSpline) = - TD.PhaseEquil_pTq(thermo_params, p, T(z), q_tot(z)) + TD.PhaseEquil_pTq( + thermo_params, + p, + oftype(p, T(z)), + oftype(p, q_tot(z)), + ) ts(p, z, ::Nothing, θ::FunctionOrSpline, q_tot::FunctionOrSpline) = - TD.PhaseEquil_pθq(thermo_params, p, θ(z), q_tot(z)) - dp_dz(p, _, z) = - -grav * TD.air_density(thermo_params, ts(p, z, T, θ, q_tot)) + TD.PhaseEquil_pθq( + thermo_params, + p, + oftype(p, θ(z)), + oftype(p, q_tot(z)), + ) + dp_dz(p, z) = -grav * TD.air_density(thermo_params, ts(p, z, T, θ, q_tot)) return column_indefinite_integral(dp_dz, p_0, (FT(0), FT(z_max))) end