Improve design of nonorographic gravity wave parameterization #3191

Xiaoyang-Xie · 2024-07-12T20:46:57Z

Purpose

Part of #897

To-do

Remove use of getindex.
Remove all instances of Fields.bycolumn and parent .
Refactor the non_orographic_gravity_wave_forcing function to be a pointwise function.

Content

Reduce allocations.
Remove use of append!
removed many instances of Fields.bycolumn and parent.
Remove use of findlast and findfirst
Refactor non_orographic_gravity_wave_forcing to be purely functional .

I have read and checked the items on the review checklist.

src/parameterized_tendencies/gravity_wave_drag/non_orographic_gravity_wave.jl

charleskawczynski · 2024-07-13T14:29:47Z

src/parameterized_tendencies/gravity_wave_drag/non_orographic_gravity_wave.jl

+    B0 = @. sign(c_hat0) * (
+        Bw * exp(-log(2.0) * ((c * flag + c_hat0 * (1 - flag) - c0) / cw)^2) + 
+        Bn * exp(-log(2.0) * ((c * flag + c_hat0 * (1 - flag) - c0) / cn)^2)
+        )


Suggested change

B0 = @. sign(c_hat0) * (

Bw * exp(-log(2.0) * ((c * flag + c_hat0 * (1 - flag) - c0) / cw)^2) +

Bn * exp(-log(2.0) * ((c * flag + c_hat0 * (1 - flag) - c0) / cn)^2)

)

@. B0 = sign(c_hat0) * (

Bw * exp(-log(2.0) * ((c * flag + c_hat0 * (1 - flag) - c0) / cw)^2) +

Bn * exp(-log(2.0) * ((c * flag + c_hat0 * (1 - flag) - c0) / cn)^2)

)

The form B0 = @. sign(c_hat0) ... points B0 to new memory. the form I suggested will mutate B0 in-place

charleskawczynski · 2024-07-13T14:30:29Z

src/parameterized_tendencies/gravity_wave_drag/non_orographic_gravity_wave.jl

+function reduce_fun1(a, b)
+    return ifelse(a[1] < b[1], a, b)
+end
+
+function reduce_fun2(a, b)
+    return ifelse(b[1] < 0, a, b)
+end
+
+function reduce_fun3(a, b)
+    return ifelse(a[1] > 0, b, a)
+end


Suggested change

function reduce_fun1(a, b)

return ifelse(a[1] < b[1], a, b)

end

function reduce_fun2(a, b)

return ifelse(b[1] < 0, a, b)

end

function reduce_fun3(a, b)

return ifelse(a[1] > 0, b, a)

end

@inline reduce_fun1(a, b) = ifelse(a[1] < b[1], a, b)

@inline reduce_fun2(a, b) = ifelse(b[1] < 0, a, b)

@inline reduce_fun3(a, b) = ifelse(a[1] > 0, b, a)

charleskawczynski · 2024-07-13T14:32:23Z

artifacts/Artifacts.toml

+[topo-info]
+git-tree-sha1 = "9970989c13c2ad015ab5738d42fd9f5b320f1451"


Did you mean to add this to the commits?

If not, let's remove it.

charleskawczynski

I made a few suggestions, but this is a good step forward.

szy21

This looks like a good step forward. @charleskawczynski @dennisYatunin Could you take a look again?

src/parameterized_tendencies/gravity_wave_drag/non_orographic_gravity_wave.jl

szy21 · 2024-07-15T21:36:03Z

src/parameterized_tendencies/gravity_wave_drag/non_orographic_gravity_wave.jl

+@inline reduce_fun1(a, b) = ifelse(a[1] < b[1], a, b)
+@inline reduce_fun2(a, b) = ifelse(b[1] < 0, a, b)
+@inline reduce_fun3(a, b) = ifelse(a[1] > 0, b, a)


Could you come up with some more informative names for these functions?

Is these names OK?:
@inline min_distance(a, b) = ifelse(a[1] < b[1], a, b)
@inline positive_selector(a, b) = ifelse(b[1] < 0, a, b)
@inline negative_selector(a, b) = ifelse(a[1] > 0, b, a)

szy21 · 2024-07-15T22:22:21Z

src/parameterized_tendencies/gravity_wave_drag/non_orographic_gravity_wave.jl

        end
+        source_level = source_level_z.:2
+
+        Operators.column_mapreduce!(sign, +, damp_level, ᶜz)


@dennisYatunin Would something like @. damp_level = Spaces.nlevels(axes(ᶜz)) work?

I tried it and found that @. damp_level = Spaces.nlevels(axes(ᶜz)) could not work, but we can use fill!(damp_level, Spaces.nlevels(axes(ᶜz))) instead.

Yes, damp_level .= Spaces.nlevels(axes(ᶜz)) or fill!(damp_level, Spaces.nlevels(axes(ᶜz)) would be better here

szy21 · 2024-07-15T22:24:35Z

src/parameterized_tendencies/gravity_wave_drag/non_orographic_gravity_wave.jl

@@ -388,3 +446,7 @@ end
 function calc_intermitency(ρ_source_level, source_ampl, nk, Bsum)
    return (source_ampl / ρ_source_level / nk) / Bsum
 end
+
+@inline reduce_fun1(a, b) = ifelse(a[1] < b[1], a, b)
+@inline reduce_fun2(a, b) = ifelse(b[1] < 0, a, b)


Suggested change

@inline reduce_fun2(a, b) = ifelse(b[1] < 0, a, b)

@inline reduce_fun2(a, b) = ifelse(b[1] <= 0, a, b)

(I think this is exactly the same as the previous logic, although it probably doesn't matter much)

szy21 · 2024-07-15T22:24:47Z

src/parameterized_tendencies/gravity_wave_drag/non_orographic_gravity_wave.jl

+
+@inline reduce_fun1(a, b) = ifelse(a[1] < b[1], a, b)
+@inline reduce_fun2(a, b) = ifelse(b[1] < 0, a, b)
+@inline reduce_fun3(a, b) = ifelse(a[1] > 0, b, a)


Suggested change

@inline reduce_fun3(a, b) = ifelse(a[1] > 0, b, a)

@inline reduce_fun3(a, b) = ifelse(a[1] >= 0, b, a)

charleskawczynski · 2024-07-16T00:36:05Z

src/parameterized_tendencies/gravity_wave_drag/non_orographic_gravity_wave.jl

-    append!(ᶜρ, ᶜρ[end] * ᶜρ[end] / ᶜρ[end - 1])
-    append!(ᶜbf, ᶜbf[end])
-    append!(ᶜz, FT(2) * ᶜz[end] - ᶜz[end - 1])
+    ᶜu = vcat(old_ᶜu, FT(2) * old_ᶜu[end] - old_ᶜu[end - 1])


Just to note, vcat is not really fixing the issue here, but the change is fine.

szy21 · 2024-07-18T21:40:06Z

Please feel free to rebase and merge, thanks!

szy21 requested review from charleskawczynski and dennisYatunin July 12, 2024 20:54

charleskawczynski reviewed Jul 13, 2024

View reviewed changes

src/parameterized_tendencies/gravity_wave_drag/non_orographic_gravity_wave.jl Outdated Show resolved Hide resolved

charleskawczynski reviewed Jul 13, 2024

View reviewed changes

src/parameterized_tendencies/gravity_wave_drag/non_orographic_gravity_wave.jl Outdated Show resolved Hide resolved

charleskawczynski reviewed Jul 13, 2024

View reviewed changes

src/parameterized_tendencies/gravity_wave_drag/non_orographic_gravity_wave.jl Show resolved Hide resolved

charleskawczynski reviewed Jul 13, 2024

View reviewed changes

Xiaoyang-Xie force-pushed the xxy/gravity_waves branch 4 times, most recently from 3d65dc9 to 2d0612e Compare July 15, 2024 21:09

refactor gravity wave cache and tendency

d688218

Xiaoyang-Xie force-pushed the xxy/gravity_waves branch from 2d0612e to d688218 Compare July 15, 2024 21:11

szy21 approved these changes Jul 15, 2024

View reviewed changes

charleskawczynski reviewed Jul 16, 2024

View reviewed changes

change function names

e2b3287

Xiaoyang-Xie added this pull request to the merge queue Jul 18, 2024

Merged via the queue into main with commit b476f7d Jul 19, 2024
11 of 14 checks passed

Xiaoyang-Xie deleted the xxy/gravity_waves branch July 19, 2024 02:34

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Improve design of nonorographic gravity wave parameterization #3191

Improve design of nonorographic gravity wave parameterization #3191

Xiaoyang-Xie commented Jul 12, 2024 •

edited by szy21

Loading

charleskawczynski Jul 13, 2024

charleskawczynski Jul 13, 2024

charleskawczynski Jul 13, 2024

charleskawczynski Jul 13, 2024

charleskawczynski left a comment

szy21 left a comment

szy21 Jul 15, 2024

Xiaoyang-Xie Jul 15, 2024

szy21 Jul 15, 2024

Xiaoyang-Xie Jul 15, 2024

dennisYatunin Jul 16, 2024

szy21 Jul 15, 2024

szy21 Jul 15, 2024

charleskawczynski Jul 16, 2024

szy21 commented Jul 18, 2024

		[topo-info]
		git-tree-sha1 = "9970989c13c2ad015ab5738d42fd9f5b320f1451"

	@inline reduce_fun2(a, b) = ifelse(b[1] < 0, a, b)
	@inline reduce_fun2(a, b) = ifelse(b[1] <= 0, a, b)

	@inline reduce_fun3(a, b) = ifelse(a[1] > 0, b, a)
	@inline reduce_fun3(a, b) = ifelse(a[1] >= 0, b, a)

Improve design of nonorographic gravity wave parameterization #3191

Improve design of nonorographic gravity wave parameterization #3191

Conversation

Xiaoyang-Xie commented Jul 12, 2024 • edited by szy21 Loading

Purpose

To-do

Content

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

charleskawczynski left a comment

Choose a reason for hiding this comment

szy21 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

szy21 commented Jul 18, 2024

Xiaoyang-Xie commented Jul 12, 2024 •

edited by szy21

Loading