Skip to content

Commit

Permalink
Fixes
Browse files Browse the repository at this point in the history
  • Loading branch information
charleskawczynski committed Oct 3, 2023
1 parent a5f7490 commit 283f32e
Showing 1 changed file with 38 additions and 17 deletions.
55 changes: 38 additions & 17 deletions src/initial_conditions/initial_conditions.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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}
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit 283f32e

Please sign in to comment.