From 283f32e677d9895451c556bcb1191f22780265f2 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 | 55 ++++++++++++++------
 1 file changed, 38 insertions(+), 17 deletions(-)

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