diff --git a/src/initial_conditions/initial_conditions.jl b/src/initial_conditions/initial_conditions.jl index 2b4b31e39a2..52e450b430e 100644 --- a/src/initial_conditions/initial_conditions.jl +++ b/src/initial_conditions/initial_conditions.jl @@ -20,70 +20,82 @@ 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: +Solve for ϕ: ``` -∂x/∂z = f(x,z) -(ᶠ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) +∂ϕ/∂z = f(ϕ,z) +(ᶠϕ^{k+1}-ᶠϕ^{k})/ᶜΔz = ᶜf(ᶜϕ̄,ᶜz) +ᶜϕ̄ = (ϕ^{k+1}+ϕ^{k})/2 +(ᶠϕ^{k+1}-ᶠϕ^{k})/ᶜΔz = ᶜf((ᶠϕ^{k+1}+ᶠϕ^{k})/2,ᶜz) +root equation: (_ϕ-ϕ^{k})/Δz = f((_ϕ+ϕ^{k})/2,ᶜz) ``` =# import RootSolvers +import ClimaComms +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 +import ClimaCore.Utilities: half function column_indefinite_integral!( f::Function, ᶠintegral::Fields.ColumnField, - x₀, + ϕ₀, ᶜzfield::Fields.ColumnField, average = (a, b) -> (a + b) / 2, ) face_space = axes(ᶠintegral) first_level = Operators.left_idx(face_space) last_level = Operators.right_idx(face_space) - ᶜΔzfield = Fields.Δz_field(ᶜz) - @inbounds Fields.level(ᶠintegral, first_level)[] = x₀ - x₁ = x₀ + ᶜΔzfield = Fields.Δz_field(ᶜzfield) + @inbounds Fields.level(ᶠintegral, first_level)[] = ϕ₀ + ϕ₁ = ϕ₀ 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 - ϕ₀) / ᶜΔz - f(average(_x, ϕ₀), ᶜz) end sol = RootSolvers.find_zero( root_eq, - RootSolvers.NewtonsMethodAD(x₀), + RootSolvers.NewtonsMethodAD(ϕ₀), RootSolvers.CompactSolution(), ) - x₁ = sol.root - f₁ = f(average(x₀, x₁), ᶜz) + ϕ₁ = sol.root + f₁ = f(average(ϕ₀, ϕ₁), ᶜ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 """ - column_indefinite_integral(dp_dz, p_0, (FT(0), z_max)) + column_indefinite_integral(f, ϕ₀, zspan; nelems = 100) The column integral, returned as an interpolate-able field. """ function column_indefinite_integral( f::Function, - x₀, + ϕ₀::FT, zspan::Tuple{FT, FT}; nelems = 100, # sets resolution for integration ) where {FT <: Real} @@ -101,7 +113,7 @@ function column_indefinite_integral( # --- zc = Fields.coordinate_field(cspace) ᶠintegral = Fields.Field(FT, fspace) - column_indefinite_integral!(f, ᶠintegral, x₀, zc) + column_indefinite_integral!(f, ᶠintegral, ϕ₀, zc) return ColumnInterpolatableField(ᶠintegral) end @@ -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