Add kinetic energy

atmos_dyamond_summer/create_artifact.jl

@@ -23,25 +23,27 @@ Plvl(z) = P0 * exp(-z / H_EARTH)
 Plvl_inv(P) = -H_EARTH * log(P / P0)
 Plvl_inv2(log_p_by_p0) = -H_EARTH * log_p_by_p0
 
-log_P(hyam, hybm, surfacepress) = log(hyam + hybm * surfacepress)
+compute_P(hyam, hybm, surfacepress) = hyam + hybm * surfacepress
 
 param_set = TP.ThermodynamicsParameters(FT)
 
 function interpz_3d(target_z, z3d, var3d)
     nx, ny, nz = size(z3d)
     target_var3d = similar(var3d, nx, ny, length(target_z))
-    for yi in 1:ny
-        for xi in 1:nx
-            source_z = view(z3d, xi, yi, :)
-            itp = extrapolate(
-                interpolate(
-                    (source_z,),
-                    (view(var3d, xi, yi, :)),
-                    Gridded(Linear()),
-                ),
-                Flat(),
-            )
-            target_var3d[xi, yi, :] .= itp.(target_z)
+    @inbounds begin
+        for yi in 1:ny
+            for xi in 1:nx
+                source_z = view(z3d, xi, yi, :)
+                itp = extrapolate(
+                    interpolate(
+                        (source_z,),
+                        (view(var3d, xi, yi, :)),
+                        Gridded(Linear()),
+                    ),
+                    Flat(),
+                )
+                target_var3d[xi, yi, :] .= itp.(target_z)
+            end
         end
     end
 
@@ -65,10 +67,11 @@ function create_initial_conditions(infile, outfile; skip = 1, small=false)
     hyam = repeat(reshape(ncin["hyam"][zidx], 1, 1, source_nz), nlon, nlat, 1)
     hybm = repeat(reshape(ncin["hybm"][zidx], 1, 1, source_nz), nlon, nlat, 1)
 
-    log_p_by_p0 = log_P.(hyam, hybm, sp) .- log_P0
+    P = FT.(compute_P.(hyam, hybm, sp))
+    log_p_by_p0 = log.(P) .- log_P0
     source_z = Plvl_inv2.(log_p_by_p0)
 
-    nz = 20 #length(zidx)
+    nz = source_nz #length(zidx)
 
     # create a target z grid with nz points
     target_z = Plvl_inv.(range(Plvl(z_min), Plvl(z_max), nz))
@@ -94,6 +97,17 @@ function create_initial_conditions(infile, outfile; skip = 1, small=false)
         )
         z[:] = target_z
 
+        # p (pressure)
+        p = defVar(ncout, "p", FT, ("lon", "lat", "z",), attrib = OrderedDict(
+            "Datatype" => string(FT),
+            "standard_name" => "pressure",
+            "long_name" => "pressure",
+            "units" => "Pa",
+            ),
+        )
+        p[:, :, :] = P
+
+
         # u (eastward_wind)
         u = defVar(ncout, "u", FT, ("lon", "lat", "z",), attrib = ncin["u"].attrib)
         u[:, :, :] = interpz_3d(target_z, source_z, ncin["u"][lonidx, latidx, zidx, 1])
@@ -112,7 +126,8 @@ function create_initial_conditions(infile, outfile; skip = 1, small=false)
 
         # q (specific_humidity)
         q = defVar(ncout, "q", FT, ("lon", "lat", "z",), attrib = ncin["q"].attrib)
-        q[:, :, :] = interpz_3d(target_z, source_z, ncin["q"][lonidx, latidx, zidx, 1])
+        _q = interpz_3d(target_z, source_z, ncin["q"][lonidx, latidx, zidx, 1])
+        q[:, :, :] = max.(_q, FT(0))
 
         # clwc (Specific cloud liquid water content)
         clwc = defVar(ncout, "clwc", FT, ("lon", "lat", "z",), attrib = ncin["clwc"].attrib)
@@ -129,11 +144,27 @@ function create_initial_conditions(infile, outfile; skip = 1, small=false)
         # cswc (Specific snow water content)
         cswc = defVar(ncout, "cswc", FT, ("lon", "lat", "z",), attrib = ncin["cswc"].attrib)
         cswc[:, :, :] = interpz_3d(target_z, source_z, ncin["cswc"][lonidx, latidx, zidx, 1])
+
+        phase_partition = TD.PhasePartition.(q)
+        ρ = TD.air_density.(param_set, t[:, :, :], P, phase_partition)
+
+        # rho (density)
+        rho = defVar(ncout, "rho", FT, ("lon", "lat", "z",), attrib = OrderedDict(
+            "Datatype" => string(FT),
+            "standard_name" => "air_density",
+            "long_name" => "air_density",
+            "units" => "Kg/m**3",
+            ),
+        )
+        rho[:, :, :] = ρ
+
+        # total energy
+        e_kin = (u .* u .+ v .* v + w .* w) .* 0.5
     end
     close(ncout)
     return nothing
 end
 #println("creating initial conditions for 0.14⁰ resolution")
 #create_initial_conditions(FILE_PATH, "DYAMOND_SUMMER_ICS_p14deg.nc")
 println("creating initial conditions for 0.98⁰ resolution")
-create_initial_conditions(FILE_PATH, "DYAMOND_SUMMER_ICS_p98deg.nc", skip=7)
+@time create_initial_conditions(FILE_PATH, "DYAMOND_SUMMER_ICS_p98deg.nc", skip=7)