diff --git a/src/cache/diagnostic_edmf_precomputed_quantities.jl b/src/cache/diagnostic_edmf_precomputed_quantities.jl
index 03a8e8a47b..8573f0ecca 100644
--- a/src/cache/diagnostic_edmf_precomputed_quantities.jl
+++ b/src/cache/diagnostic_edmf_precomputed_quantities.jl
@@ -371,13 +371,7 @@ function set_diagnostic_edmf_precomputed_quantities_do_integral!(Y, p, t)
 
             # We don't have an upper limit to entrainment for the first level 
             # (calculated at i=2), as the vertical at the first level is zero
-            if i == 2
-                @. entrʲ_prev_level = limit_entrainment(
-                    entrʲ_prev_level,
-                    draft_area(ρaʲ_prev_level, ρʲ_prev_level),
-                    dt,
-                )
-            else
+            if i > 2
                 @. entrʲ_prev_level = limit_entrainment(
                     entrʲ_prev_level,
                     draft_area(ρaʲ_prev_level, ρʲ_prev_level),
@@ -388,6 +382,11 @@ function set_diagnostic_edmf_precomputed_quantities_do_integral!(Y, p, t)
                     dz_prev_level,
                 )
             end
+            @. entrʲ_prev_level = limit_entrainment(
+                entrʲ_prev_level,
+                draft_area(ρaʲ_prev_level, ρʲ_prev_level),
+                dt,
+            )
 
             # TODO: use updraft top instead of scale height
             @. nh_pressure³ʲ_prev_halflevel = ᶠupdraft_nh_pressure(
@@ -517,6 +516,15 @@ function set_diagnostic_edmf_precomputed_quantities_do_integral!(Y, p, t)
                 p.atmos.edmfx_detr_model,
             )
 
+            @. detrʲ_prev_level = limit_detrainment(
+                detrʲ_prev_level,
+                draft_area(ρaʲ_prev_level, ρʲ_prev_level),
+                get_physical_w(
+                    u³ʲ_prev_halflevel,
+                    local_geometry_prev_halflevel,
+                ),
+                dz_prev_level,
+            )
             @. detrʲ_prev_level = limit_detrainment(
                 detrʲ_prev_level,
                 draft_area(ρaʲ_prev_level, ρʲ_prev_level),
diff --git a/src/prognostic_equations/edmfx_tke.jl b/src/prognostic_equations/edmfx_tke.jl
index f30c893e27..ff8a843a28 100644
--- a/src/prognostic_equations/edmfx_tke.jl
+++ b/src/prognostic_equations/edmfx_tke.jl
@@ -50,11 +50,14 @@ function edmfx_tke_tendency!(
         # entrainment and detraiment
         # using ᶜu⁰ and local geometry results in allocation
         for j in 1:n
+            ᶜρaʲ_colidx =
+                turbconv_model isa PrognosticEDMFX ?
+                Y.c.sgsʲs.:($j).ρa[colidx] : p.precomputed.ᶜρaʲs.:($j)[colidx]
             @. Yₜ.c.sgs⁰.ρatke[colidx] +=
-                ᶜρa⁰[colidx] * (
-                    ᶜentrʲs.:($$j)[colidx] * 1 / 2 * norm_sqr(
+                ᶜρaʲ_colidx * (
+                    ᶜdetrʲs.:($$j)[colidx] * 1 / 2 * norm_sqr(
                         ᶜinterp(ᶠu³⁰[colidx]) - ᶜinterp(ᶠu³ʲs.:($$j)[colidx]),
-                    ) - ᶜdetrʲs.:($$j)[colidx] * ᶜtke⁰[colidx]
+                    ) - ᶜentrʲs.:($$j)[colidx] * ᶜtke⁰[colidx]
                 )
         end
         # pressure work