From a4e46a152e1e649832124fd350cf036485a31913 Mon Sep 17 00:00:00 2001
From: Charles Kawczynski <kawczynski.charles@gmail.com>
Date: Mon, 2 Oct 2023 13:09:12 -0700
Subject: [PATCH] Fixes

---
 src/initial_conditions/initial_conditions.jl | 81 ++++++++++++--------
 1 file changed, 51 insertions(+), 30 deletions(-)

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