From 0c2adfa88461cf9223a541fa71ed4412c6993e24 Mon Sep 17 00:00:00 2001
From: Dave Rumph <daverumph@icloud.com>
Date: Tue, 25 Nov 2025 16:04:09 -0800
Subject: [PATCH] Rewrite update_cloud_diagnostics as a point function to avoid
 GPU compiler issues

A rewrite of ClimaCore.Broadcast.getindex to address other GPU kernel issues broke this code.
---
 src/cache/cloud_fraction.jl | 37 ++++++++++++++++++-------------------
 1 file changed, 18 insertions(+), 19 deletions(-)

diff --git a/src/cache/cloud_fraction.jl b/src/cache/cloud_fraction.jl
index e8d16b1632..5114916a4f 100644
--- a/src/cache/cloud_fraction.jl
+++ b/src/cache/cloud_fraction.jl
@@ -123,6 +123,23 @@ NVTX.@annotate function set_cloud_fraction!(
 end
 
 
+function updraft_cloud_diagnostics(ρʲs, ρaʲs, tsʲs, thermo_params, n)
+    FT = eltype(thermo_params)
+    cf = FT(0)
+    q_liq = FT(0)
+    q_ice = FT(0)
+    for j in 1:n
+        aʲ = draft_area(ρaʲs.:($j), ρʲs.:($j))
+        if TD.has_condensate(thermo_params, tsʲs.:($j))
+            cf += aʲ
+        end
+        (; liq, ice) = TD.PhasePartition(thermo_params, tsʲs.:($j))
+        q_liq += aʲ * liq
+        q_ice += aʲ * ice
+    end
+    return (; cf, q_liq, q_ice)
+end
+
 NVTX.@annotate function set_cloud_fraction!(
     Y,
     p,
@@ -159,24 +176,8 @@ NVTX.@annotate function set_cloud_fraction!(
     n = n_mass_flux_subdomains(turbconv_model)
     if n > 0
         (; ᶜρaʲs, ᶜρʲs, ᶜtsʲs) = p.precomputed
+        @. cloud_diagnostics_tuple += updraft_cloud_diagnostics(ᶜρaʲs, ᶜρʲs, ᶜtsʲs, thermo_params, n)
     end
-
-    for j in 1:n
-        @. cloud_diagnostics_tuple += NamedTuple{(:cf, :q_liq, :q_ice)}(
-            tuple(
-                ifelse(
-                    TD.has_condensate(thermo_params, ᶜtsʲs.:($$j)),
-                    draft_area(ᶜρaʲs.:($$j), ᶜρʲs.:($$j)),
-                    0,
-                ),
-                draft_area(ᶜρaʲs.:($$j), ᶜρʲs.:($$j)) *
-                TD.PhasePartition(thermo_params, ᶜtsʲs.:($$j)).liq,
-                draft_area(ᶜρaʲs.:($$j), ᶜρʲs.:($$j)) *
-                TD.PhasePartition(thermo_params, ᶜtsʲs.:($$j)).ice,
-            ),
-        )
-    end
-
 end
 
 NVTX.@annotate function set_cloud_fraction!(
@@ -309,8 +310,6 @@ function quad_loop(
 
     function f(x1_hat, x2_hat)
         FT = eltype(x1_hat)
-        @assert(x1_hat >= FT(0))
-        @assert(x2_hat >= FT(0))
         _ts = thermo_state(thermo_params; p = p_c, θ = x1_hat, q_tot = x2_hat)
         hc = TD.has_condensate(thermo_params, _ts)