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 ae3fd21
Showing 1 changed file with 52 additions and 31 deletions.
83 changes: 52 additions & 31 deletions src/initial_conditions/initial_conditions.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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,
average = (ϕ⁺, ϕ⁻) -> (ϕ⁺ + ϕ⁻) / 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}
Expand All @@ -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

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 ae3fd21

Please sign in to comment.