From 3b6648b2b491a5489d9d77df22906fa8add2d37c Mon Sep 17 00:00:00 2001
From: Anna Jaruga <ajaruga@caltech.edu>
Date: Wed, 20 Nov 2024 16:24:47 -0800
Subject: [PATCH] Add perf and GPU tests for Chen2022 terminal velocity

---
 test/gpu_tests.jl         | 60 +++++++++++++++++++++++++++++++++++++++
 test/performance_tests.jl | 15 ++++++++++
 2 files changed, 75 insertions(+)

diff --git a/test/gpu_tests.jl b/test/gpu_tests.jl
index cde5b1425..182f8db82 100644
--- a/test/gpu_tests.jl
+++ b/test/gpu_tests.jl
@@ -131,6 +131,32 @@ end
     end
 end
 
+@kernel function test_chen2022_terminal_velocity_kernel!(
+    liquid,
+    ice,
+    rain,
+    snow,
+    Ch2022,
+    output::AbstractArray{FT},
+    ρₐ,
+    qₗ,
+    qᵢ,
+    qᵣ,
+    qₛ,
+) where {FT}
+
+    i = @index(Group, Linear)
+
+    @inbounds begin
+        output[1, i] = CMN.terminal_velocity(liquid, Ch2022.rain, ρₐ[i], qₗ[i])
+        output[2, i] =
+            CMN.terminal_velocity(ice, Ch2022.small_ice, ρₐ[i], qᵢ[i])
+        output[3, i] = CM1.terminal_velocity(rain, Ch2022.rain, ρₐ[i], qᵣ[i])
+        output[4, i] =
+            CM1.terminal_velocity(snow, Ch2022.large_ice, ρₐ[i], qₛ[i])
+    end
+end
+
 @kernel function test_0_moment_micro_kernel!(
     p0m,
     output::AbstractArray{FT},
@@ -750,6 +776,7 @@ function test_gpu(FT)
     # Terminal velocity
     blk1mvel = CMP.Blk1MVelType(FT)
     SB2006Vel = CMP.SB2006VelType(FT)
+    Ch2022 = CMP.Chen2022VelType(FT)
 
     # 2-moment microphysics
     SB2006 = CMP.SB2006(FT)
@@ -829,6 +856,39 @@ function test_gpu(FT)
         @test Array(output)[2] ≈ FT(-1e-4)
     end
 
+    @testset "Chen 2022 terminal velocity kernels" begin
+        dims = (4, 1)
+        (; output, ndrange) = setup_output(dims, FT)
+
+        ρ = ArrayType([FT(0.95)])
+        qₗ = ArrayType([FT(0.004)])
+        qᵢ = ArrayType([FT(0.003)])
+        qᵣ = ArrayType([FT(0.002)])
+        qₛ = ArrayType([FT(0.001)])
+
+        kernel! = test_chen2022_terminal_velocity_kernel!(backend, work_groups)
+        kernel!(
+            liquid,
+            ice,
+            rain,
+            snow,
+            Ch2022,
+            output,
+            ρ,
+            qₗ,
+            qᵢ,
+            qᵣ,
+            qₛ;
+            ndrange,
+        )
+
+        # test that terminal velocity is callable and returns a reasonable value
+        @test Array(output)[1] ≈ FT(0.00689286343412659)
+        @test Array(output)[2] ≈ FT(0.011493257177438487)
+        @test Array(output)[3] ≈ FT(4.274715870117866)
+        @test Array(output)[4] ≈ FT(0.5688979454130587)
+    end
+
     @testset "0-moment microphysics kernels" begin
         dims = (3, 2)
         (; output, ndrange) = setup_output(dims, FT)
diff --git a/test/performance_tests.jl b/test/performance_tests.jl
index 964fe9c3f..e0fdf1f39 100644
--- a/test/performance_tests.jl
+++ b/test/performance_tests.jl
@@ -66,6 +66,7 @@ function benchmark_test(FT)
     liquid = CMP.CloudLiquid(FT)
     ice = CMP.CloudIce(FT)
     rain = CMP.Rain(FT)
+    snow = CMP.Snow(FT)
     ce = CMP.CollisionEff(FT)
     # 2-moment microphysics
     override_file = joinpath(
@@ -202,6 +203,20 @@ function benchmark_test(FT)
     #bench_press(CMI_het.P3_deposition_N_i, (ip.p3, T_air_cold), 230)
     #bench_press(CMI_het.P3_het_N_i, (ip.p3, T_air_cold, N_liq, V_liq, Δt), 230)
 
+    # Chen 2022 terminal velocity
+    bench_press(CMN.terminal_velocity, (liquid, ch2022.rain, ρ_air, q_liq), 350)
+    bench_press(
+        CMN.terminal_velocity,
+        (ice, ch2022.small_ice, ρ_air, q_ice),
+        350,
+    )
+    bench_press(CM1.terminal_velocity, (rain, ch2022.rain, ρ_air, q_rai), 700)
+    bench_press(
+        CM1.terminal_velocity,
+        (snow, ch2022.large_ice, ρ_air, q_sno),
+        700,
+    )
+
     # aerosol activation
     bench_press(
         AA.total_N_activated,