From ff07b43fd6ed3c247915864b19d48ddb0fd192ce Mon Sep 17 00:00:00 2001
From: Zhaoyi Shen <11598433+szy21@users.noreply.github.com>
Date: Wed, 13 Mar 2024 14:42:21 -0700
Subject: [PATCH] refactor surface albedo calculation

---
 src/cache/surface_albedo.jl | 37 ++++++++++++++++---------------------
 1 file changed, 16 insertions(+), 21 deletions(-)

diff --git a/src/cache/surface_albedo.jl b/src/cache/surface_albedo.jl
index 1d579603ec..faaba37e14 100644
--- a/src/cache/surface_albedo.jl
+++ b/src/cache/surface_albedo.jl
@@ -80,29 +80,24 @@ function set_surface_albedo!(
     λ = FT(0) # spectral wavelength (not used for now)
     μ = cos_zenith
 
-    ArrayType = ClimaComms.array_type(Y)
-
-    # set the direct and diffuse surface albedos
-    direct_sw_surface_albedo .=
-        ArrayType(ones(axes(direct_sw_surface_albedo))) .* ArrayType(
-            reshape(
-                surface_albedo_direct(
-                    α_model,
-                ).(λ, μ, RRTMGPI.field2array(norm.(Fields.level(Y.c.uₕ, 1)))),
-                1,
-                :,
-            ),
+    surface_albedo = p.scratch.temp_field_level
+    f_direct = surface_albedo_direct(α_model)
+    surface_albedo .=
+        f_direct.(
+            λ,
+            RRTMGPI.array2field(μ, axes(surface_albedo)),
+            norm.(Fields.level(Y.c.uₕ, 1)),
         )
-    diffuse_sw_surface_albedo .=
-        ArrayType(ones(axes(direct_sw_surface_albedo))) .* ArrayType(
-            reshape(
-                surface_albedo_diffuse(
-                    α_model,
-                ).(λ, μ, RRTMGPI.field2array(norm.(Fields.level(Y.c.uₕ, 1)))),
-                1,
-                :,
-            ),
+    direct_sw_surface_albedo .= RRTMGPI.field2array(surface_albedo)'
+
+    f_diffuse = surface_albedo_diffuse(α_model)
+    surface_albedo .=
+        f_diffuse.(
+            λ,
+            RRTMGPI.array2field(μ, axes(surface_albedo)),
+            norm.(Fields.level(Y.c.uₕ, 1)),
         )
+    diffuse_sw_surface_albedo .= RRTMGPI.field2array(surface_albedo)'
 end
 
 """