try fixing GPU runs

src/cache/cloud_fraction.jl

@@ -61,12 +61,12 @@ function set_cloud_fraction!(Y, p, ::Union{EquilMoistModel, NonEquilMoistModel})
 end
 
 """
-    function quad_loop(SG_quad, p_c, qt_mean, θl_mean, qt′qt′, θl′θl′, θl′qt′, thermo_params)
+    function quad_loop(SG_quad, p_c, q_mean, θ_mean, qt′qt′, θl′θl′, θl′qt′, thermo_params)
 
 where:
   - SG_quad is a struct containing information about quadrature type and order
   - p_c is the atmospheic pressure
-  - qt_mean, θl_mean is the grid mean q_tot and liquid ice potential temperature
+  - q_mean, θ_mean is the grid mean q_tot and liquid ice potential temperature
   - qt′qt′, θl′θl′, θl′qt′ are the (co)variances of q_tot and liquid ice potential temperature
   - thermo params are the thermodynamics parameters
 
@@ -76,82 +76,91 @@ It returns cloud fraction computed as a sum over quadrature points.
 function quad_loop(
     SG_quad::SGSQuadrature,
     p_c,
-    qt_mean,
-    θl_mean,
-    qt′qt′,
-    θl′θl′,
-    θl′qt′,
+    q_mean,
+    θ_mean,
+    q′q′,
+    θ′θ′,
+    θ′q′,
     thermo_params,
 )
-    FT = eltype(qt_mean)
+    FT = eltype(q_mean)
 
     @assert SG_quad.quadrature_type isa GaussianQuad
 
     sqrt2 = FT(sqrt(2))
 
     # Epsilon defined per typical variable fluctuation
-    eps_q = FT(eps(FT)) * max(FT(eps(FT)), qt_mean)
-    eps_θ = FT(eps(FT))
+    eps_q = eps(FT) * max(eps(FT), q_mean)
+    eps_θ = eps(FT)
 
     # limit σ_q to prevent negative q_tot_hat
-    σ_q_lim = -qt_mean / (sqrt2 * SG_quad.a[1])
-    σ_q::FT = min(sqrt(qt′qt′), σ_q_lim)
-    σ_θ::FT = sqrt(θl′θl′)
+    σ_q_lim = -q_mean / (sqrt2 * SG_quad.a[1])
+    σ_q = min(sqrt(q′q′), σ_q_lim)
+    # Do we also have to try to limit θ in the same way as q??
+    σ_θ = sqrt(θ′θ′)
 
     # Enforce Cauchy-Schwarz inequality, numerically stable compute
-    _corr::FT = (θl′qt′ / max(σ_q, eps_q))
-    corr::FT = max(min(_corr / max(σ_θ, eps_θ), 1), -1)
+    _corr = (θ′q′ / max(σ_q, eps_q))
+    corr = max(min(_corr / max(σ_θ, eps_θ), FT(1)), FT(-1))
 
     # Conditionals
     σ_c = sqrt(max(1 - corr * corr, 0)) * σ_θ
 
-    function get_x_hat(χ::Tuple{<:Real, <:Real})
-        μ_c = θl_mean + sqrt2 * corr * σ_θ * χ[1]
-        θ_hat = μ_c + sqrt2 * σ_c * χ[2]
-        q_tot_hat = qt_mean + sqrt2 * σ_q * χ[1]
+    # Returns the physical values based on quadrature sampling points
+    # and limited covarainces
+    function get_x_hat(χ1, χ2)
+        FT = eltype(χ1)
+        μ_c = θ_mean + sqrt2 * corr * σ_θ * χ1
+        θ_hat = μ_c + sqrt2 * σ_c * χ2
+        q_hat = q_mean + sqrt2 * σ_q * χ1
         # The σ_q_lim limits q_tot_hat to be close to zero
         # for the negative sampling points. However due to numerical erros
         # we sometimes still get small negative numers here
-        return (θ_hat, max(FT(0), q_tot_hat))
+        return (θ_hat, max(FT(0), q_hat))
     end
 
     # cloudy/dry categories for buoyancy in TKE
-    f_q_tot_sat(x_hat::Tuple{<:Real, <:Real}, hc) =
-        hc ? x_hat[2] : eltype(x_hat)(0)
-
-    get_ts(x_hat::Tuple{<:Real, <:Real}) =
-        thermo_state(thermo_params; p = p_c, θ = x_hat[1], q_tot = x_hat[2])
-    f_cf(x_hat::Tuple{<:Real, <:Real}, hc) =
-        hc ? eltype(x_hat)(1) : eltype(x_hat)(0)
-    function f(x_hat::Tuple{<:Real, <:Real})
-        @assert(x_hat[1] >= FT(0))
-        @assert(x_hat[2] >= FT(0))
-        ts = get_ts(x_hat)
+    f_q_tot_sat(x1_hat, x2_hat, hc) =
+        hc ? x2_hat : eltype(x2_hat)(0)
+    get_ts(x1_hat, x2_hat) =
+        thermo_state(thermo_params; p = p_c, θ = x1_hat, q_tot = x2_hat)
+    f_cf(x1_hat, x2_hat, hc) =
+        hc ? eltype(x1_hat)(1) : eltype(x1_hat)(0)
+    function f(x1_hat, x2_hat)
+        @assert(x1_hat >= FT(0))
+        @assert(x2_hat >= FT(0))
+        ts = get_ts(x1_hat, x2_hat)
         hc = TD.has_condensate(thermo_params, ts)
-        return (; cf = f_cf(x_hat, hc), q_tot_sat = f_q_tot_sat(x_hat, hc))
+        return (;
+            cf = f_cf(x1_hat, x2_hat, hc),
+            q_tot_sat = f_q_tot_sat(x1_hat, x2_hat, hc),
+        )
     end
 
     return quad(f, get_x_hat, SG_quad).cf
 end
 
 """
-   Compute f(A, B) as a sum over inner and outer quadrature points
-   that approximate the sub-grid scale variability of A and B
+   Compute f(θ, q) as a sum over inner and outer quadrature points
+   that approximate the sub-grid scale variability of θ and q.
+
+   θ - liquid ice potential temperature
+   q - total water specific humidity
 """
 function quad(f::F, get_x_hat::F1, quad) where {F <: Function, F1 <: Function}
     χ = quad.a
     weights = quad.w
     quad_order = quadrature_order(quad)
     FT = eltype(χ)
     # zero outer quadrature points
-    T = typeof(f(get_x_hat((χ[1], χ[1]))))
+    T = typeof(f(get_x_hat(χ[1], χ[1])...))
     outer_env = rzero(T)
     @inbounds for m_q in 1:quad_order
         # zero inner quadrature points
         inner_env = rzero(T)
         for m_h in 1:quad_order
-            x_hat = get_x_hat((χ[m_q], χ[m_h]))
-            inner_env = inner_env ⊞ f(x_hat) ⊠ weights[m_h] ⊠ FT(1 / sqrt(π))
+            x_hat = get_x_hat(χ[m_q], χ[m_h])
+            inner_env = inner_env ⊞ f(x_hat...) ⊠ weights[m_h] ⊠ FT(1 / sqrt(π))
         end
         outer_env = outer_env ⊞ inner_env ⊠ weights[m_q] ⊠ FT(1 / sqrt(π))
     end