FourierFlows / GeophysicalFlows.jl

Geophysical fluid dynamics pseudospectral solvers with Julia and FourierFlows.jl.
https://fourierflows.github.io/GeophysicalFlowsDocumentation/stable/
MIT License
153 stars 32 forks source link

Optimize 2-layer configs for `MultiLayerQG` #270

Closed navidcy closed 2 years ago

navidcy commented 2 years ago

This PR attempts to optimize the 2-layer configuration of MultiLayerQG module by hard-coding the PV-streamfunction relationships. To do so we create a TwoLayerParams <: AbstractParams and then use this type to add methods in both streamfunctionfrompv! and pvfromstreamfunction! functions that include the hard-coded PV-streamfunction relationships for 2-layer configurations.

Some benchmarks for 2- and 3-layer configs with nx = ny = 128:

with nlayers=2

before this PR

julia> using GeophysicalFlows, BenchmarkTools

julia> nlayers = 2; prob = MultiLayerQG.Problem(nlayers, GPU()); @btime stepforward!(prob)
3.317 s (700757 allocations: 114.42 MiB)

after this PR

julia> using GeophysicalFlows, BenchmarkTools

julia> nlayers = 2; prob = MultiLayerQG.Problem(nlayers, GPU()); @btime stepforward!(prob)
1.923 ms (3385 allocations: 321.58 KiB)

with nlayers=3 (shouldn't be affected by this PR)

before this PR

julia> using GeophysicalFlows, BenchmarkTools

julia> nlayers = 3; prob = MultiLayerQG.Problem(nlayers, GPU()); @btime stepforward!(prob)
6.194 s (1299797 allocations: 213.44 MiB)

after this PR

julia> using GeophysicalFlows, BenchmarkTools

julia> nlayers = 3; prob = MultiLayerQG.Problem(nlayers, GPU()); @btime stepforward!(prob)
5.641 s (1299797 allocations: 213.44 MiB)

1650-fold faster for 2-layer config. We'll take that as an improvement. (Possibly even faster for higher grid resolution!)

Closes #267.

navidcy commented 2 years ago

90edadb is a temporary hack

navidcy commented 2 years ago

Needs to be cleaned up. But seems to work!!

navidcy commented 2 years ago

ok @glwagner and @szy21, feel free to review this!

navidcy commented 2 years ago

Looks great to me! It seems possible to reuse one Params struct by adding a type parameter N corresponding to the number of layers. Then we can define the aliases


const SingleLayerParams = Params{1}

const TwoLayerParams = Params{2}

I don't see this as necessary for this PR though, just future cleanup.

I agree! Can you open an issue with these ideas for a future cleanup?