Skip to content

Commit

Permalink
Merge branch 'mk/sppt' of https://github.com/SpeedyWeather/SpeedyWeat…
Browse files Browse the repository at this point in the history 
…her.jl into mk/sppt
  • Loading branch information
@milankl
milankl committed Dec 5, 2024
2 parents 01d632f + 3d658aa commit b16ba40
Show file tree
Hide file tree
Showing 4 changed files with 144 additions and 105 deletions.
29 changes: 29 additions & 0 deletions .github/workflows/extended_tests.yml
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,29 @@
name: Extended Test
# Weekly extended tests
on:
schedule:
- cron: '0 2 * * 0' # Daily at 2 AM UTC every Sunday
jobs:
test:
runs-on: ${{ matrix.os }}
strategy:
matrix:
version:
- '1.11'
os:
- ubuntu-latest
arch:
- x64
steps:
- uses: actions/checkout@v4
- uses: julia-actions/setup-julia@v2
with:
version: ${{ matrix.version }}
arch: ${{ matrix.arch }}
- uses: julia-actions/cache@v2
- uses: julia-actions/julia-buildpkg@v1
- uses: julia-actions/julia-runtest@v1
with:
test_args: 'extended_tests'
env:
GITHUB_AUTH: ${{ secrets.GITHUB_TOKEN }}
1 change: 1 addition & 0 deletions CHANGELOG.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
## Unreleased

- AbstractStochasticPhysics and SPPT implemented [#629](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/629)
- Introduced seperate extended tests that are not run on every commit [#628](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/628)
- One-band longwave radiation [#624](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/624)
- compat entry for FiniteDifferences.jl [#620](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/620)
- Slab ocean and net surface fluxes [#613](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/613)
Expand Down
6 changes: 6 additions & 0 deletions test/runtests.jl
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,12 @@
using SpeedyWeather
using Test

FLAG_EXTENDED_TESTS = "extended_tests" in ARGS ? true : false

if FLAG_EXTENDED_TESTS
@info "Running extended test suite"
end

# GENERAL
include("utility_functions.jl")
include("dates.jl")
Expand Down
213 changes: 108 additions & 105 deletions test/spectral_transform_ad_rules.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ import AbstractFFTs

grid_types = [FullGaussianGrid, OctahedralGaussianGrid] # one full and one reduced grid, both Gaussian to have exact transforms
grid_dealiasing = [2, 3]
fd_tests = [true, false] # to save CI time, only do FiniteDifferences test for one of the grids
fd_tests = [true, true]

# currenlty there's an issue with EnzymeTestUtils not being able to work with structs with undefined fields like FFT plans
# https://github.com/EnzymeAD/Enzyme.jl/issues/1992
Expand Down Expand Up @@ -56,131 +56,134 @@ end
end
end

@testset "Complete Transform Enzyme" begin
# make a high level finite difference test of the whole transform
# can't use Enzyme or ChainRule Test tools for tests for that
for (i_grid, grid_type) in enumerate(grid_types)
if FLAG_EXTENDED_TESTS # part of the extented tests, not tested in regular CI

spectral_grid = SpectralGrid(Grid=grid_type, trunc=10, nlayers=1, dealiasing=grid_dealiasing[i_grid])
S = SpectralTransform(spectral_grid, one_more_degree=true)
dS = deepcopy(S)
@testset "Complete Transform Enzyme" begin
# make a high level finite difference test of the whole transform
# can't use Enzyme or ChainRule Test tools for tests for that
for (i_grid, grid_type) in enumerate(grid_types)

if fd_tests[i_grid]

# forwards
grid = rand(spectral_grid.Grid{spectral_grid.NF}, spectral_grid.nlat_half, spectral_grid.nlayers)
dgrid = zero(grid)
specs = zeros(LowerTriangularArray{Complex{spectral_grid.NF}}, spectral_grid.trunc+2, spectral_grid.trunc+1, spectral_grid.nlayers)

# seed
dspecs = zero(specs)
fill!(dspecs, 1+1im)
spectral_grid = SpectralGrid(Grid=grid_type, trunc=10, nlayers=1, dealiasing=grid_dealiasing[i_grid])
S = SpectralTransform(spectral_grid, one_more_degree=true)
dS = deepcopy(S)

autodiff(Reverse, transform!, Const, Duplicated(specs, dspecs), Duplicated(grid, dgrid), Duplicated(S, dS))
if fd_tests[i_grid]

# forwards
grid = rand(spectral_grid.Grid{spectral_grid.NF}, spectral_grid.nlat_half, spectral_grid.nlayers)
dgrid = zero(grid)
specs = zeros(LowerTriangularArray{Complex{spectral_grid.NF}}, spectral_grid.trunc+2, spectral_grid.trunc+1, spectral_grid.nlayers)

# seed
dspecs = zero(specs)
fill!(dspecs, 1+1im)

# new seed
dspecs2 = zero(specs)
fill!(dspecs2, 1+1im)
autodiff(Reverse, transform!, Const, Duplicated(specs, dspecs), Duplicated(grid, dgrid), Duplicated(S, dS))

# finite difference comparision, seeded with a one adjoint to get the direct gradient
fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> transform(x, S), dspecs2, grid)
@test isapprox(dgrid, fd_jvp[1])
# new seed
dspecs2 = zero(specs)
fill!(dspecs2, 1+1im)

## now backwards, as the input for spec we use the output of the forward transform
# finite difference comparision, seeded with a one adjoint to get the direct gradient
fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> transform(x, S), dspecs2, grid)
@test isapprox(dgrid, fd_jvp[1])

fill!(dspecs,0)
grid = zeros(spectral_grid.Grid{spectral_grid.NF}, spectral_grid.nlat_half, spectral_grid.nlayers)
dgrid = similar(grid)
fill!(dgrid, 1)
## now backwards, as the input for spec we use the output of the forward transform

autodiff(Reverse, transform!, Const, Duplicated(grid, dgrid), Duplicated(specs, dspecs), Duplicated(S, dS))
fill!(dspecs,0)
grid = zeros(spectral_grid.Grid{spectral_grid.NF}, spectral_grid.nlat_half, spectral_grid.nlayers)
dgrid = similar(grid)
fill!(dgrid, 1)

# new seed
dgrid2 = similar(grid)
fill!(dgrid2, 1)
autodiff(Reverse, transform!, Const, Duplicated(grid, dgrid), Duplicated(specs, dspecs), Duplicated(S, dS))

fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> transform(x, S), dgrid2, specs)
# new seed
dgrid2 = similar(grid)
fill!(dgrid2, 1)

@test isapprox(dspecs, fd_jvp[1])
end
fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> transform(x, S), dgrid2, specs)

# test that d S^{-1}(S(x)) / dx = dx/dx = 1 (starting in both domains)
# this only holds for exact transforms, like Gaussian grids
@test isapprox(dspecs, fd_jvp[1])
end

# start with grid (but with a truncated one)
function transform_identity!(x_out::AbstractGridArray{T}, x::AbstractGridArray{T}, S::SpectralTransform{T}) where T
x_SH = zeros(LowerTriangularArray{Complex{T}}, S.lmax+1, S.mmax+1, S.nlayers)
transform!(x_SH, x, S)
transform!(x_out, x_SH, S)
return nothing
end
# test that d S^{-1}(S(x)) / dx = dx/dx = 1 (starting in both domains)
# this only holds for exact transforms, like Gaussian grids

function transform_identity(x::AbstractGridArray{T}, S::SpectralTransform{T}) where T
x_copy = deepcopy(x)
transform_identity!(x_copy, x, S)
return x_copy
end

grid = rand(S.Grid{spectral_grid.NF}, S.nlat_half, S.nlayers)
spec = transform(grid, S)

grid = transform(spec, S)
grid_out = zero(grid)

transform_identity!(grid_out, grid, S)
@test isapprox(grid, grid_out)

dgrid = similar(grid)
fill!(dgrid, 1)

dgrid_out = zero(grid_out)

autodiff(Reverse, transform_identity!, Const, Duplicated(grid_out, dgrid_out), Duplicated(grid, dgrid), Duplicated(S, dS))

@test all(isapprox.(dgrid, 1))
# TODO: previously this test was broken, with a version that directly mutates in-place.
# FD yields the same non-one values though.
# Not sure why. Do we use such things in our model?
#
#function transform_identity!(x::AbstractGridArray{T}, S::SpectralTransform{T}) where T
# x_SH = zeros(LowerTriangularArray{Complex{T}}, S.lmax+1, S.mmax+1, S.nlayers)
# transform!(x_SH, x, S)
# transform!(x, x_SH, S)
# return nothing
#end
# The FD comparision passes, but it takes a long time to compute, so it's commented out.
#dgrid2 = similar(grid)
#fill!(dgrid2, 1)
#fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> transform_identity(x, S), dgrid2, grid)
#@test isapprox(dgrid, fd_jvp[1], rtol=0.01)

# now start with spectral space, exclude for other grid because of https://github.com/SpeedyWeather/SpeedyWeather.jl/issues/626
if fd_tests[i_grid]

function transform_identity!(x::LowerTriangularArray{Complex{T}}, S::SpectralTransform{T}) where T
x_grid = zeros(S.Grid{T}, S.nlat_half, S.nlayers)
transform!(x_grid, x, S)
transform!(x, x_grid, S)
# start with grid (but with a truncated one)
function transform_identity!(x_out::AbstractGridArray{T}, x::AbstractGridArray{T}, S::SpectralTransform{T}) where T
x_SH = zeros(LowerTriangularArray{Complex{T}}, S.lmax+1, S.mmax+1, S.nlayers)
transform!(x_SH, x, S)
transform!(x_out, x_SH, S)
return nothing
end

function transform_identity(x::AbstractGridArray{T}, S::SpectralTransform{T}) where T
x_copy = deepcopy(x)
transform_identity!(x_copy, x, S)
return x_copy
end

grid = rand(S.Grid{spectral_grid.NF}, S.nlat_half, S.nlayers)
spec = transform(grid, S)
spec_copy = deepcopy(spec)
transform_identity!(spec, S)
@test isapprox(spec, spec_copy)

dspec = similar(spec)
fill!(dspec, 1+im)

autodiff(Reverse, transform_identity!, Const, Duplicated(spec, dspec), Duplicated(S, dS))
grid = transform(spec, S)
grid_out = zero(grid)

transform_identity!(grid_out, grid, S)
@test isapprox(grid, grid_out)

@test all(all.([isapprox.(dspec[il,1,:], 1) for il in 1:S.lmax+1])) # m = 0 => Im = 0
dgrid = similar(grid)
fill!(dgrid, 1)

for i in eachmatrix(dspec)
@test all(isapprox.(dspec[:,i][S.lmax+2:end], 1+im))
end
end
end
dgrid_out = zero(grid_out)

autodiff(Reverse, transform_identity!, Const, Duplicated(grid_out, dgrid_out), Duplicated(grid, dgrid), Duplicated(S, dS))

@test all(isapprox.(dgrid, 1))
# TODO: previously this test was broken, with a version that directly mutates in-place.
# FD yields the same non-one values though.
# Not sure why. Do we use such things in our model?
#
#function transform_identity!(x::AbstractGridArray{T}, S::SpectralTransform{T}) where T
# x_SH = zeros(LowerTriangularArray{Complex{T}}, S.lmax+1, S.mmax+1, S.nlayers)
# transform!(x_SH, x, S)
# transform!(x, x_SH, S)
# return nothing
#end
# The FD comparision passes, but it takes a long time to compute, so it's commented out.
#dgrid2 = similar(grid)
#fill!(dgrid2, 1)
#fd_jvp = FiniteDifferences.j′vp(central_fdm(5,1), x -> transform_identity(x, S), dgrid2, grid)
#@test isapprox(dgrid, fd_jvp[1], rtol=0.01)

# now start with spectral space, exclude for other grid because of https://github.com/SpeedyWeather/SpeedyWeather.jl/issues/626
if fd_tests[i_grid]

function transform_identity!(x::LowerTriangularArray{Complex{T}}, S::SpectralTransform{T}) where T
x_grid = zeros(S.Grid{T}, S.nlat_half, S.nlayers)
transform!(x_grid, x, S)
transform!(x, x_grid, S)
return nothing
end

spec = transform(grid, S)
spec_copy = deepcopy(spec)
transform_identity!(spec, S)
@test isapprox(spec, spec_copy)

dspec = similar(spec)
fill!(dspec, 1+im)

autodiff(Reverse, transform_identity!, Const, Duplicated(spec, dspec), Duplicated(S, dS))

@test all(all.([isapprox.(dspec[il,1,:], 1) for il in 1:S.lmax+1])) # m = 0 => Im = 0

for i in eachmatrix(dspec)
@test all(isapprox.(dspec[:,i][S.lmax+2:end], 1+im))
end
end
end
end
end

@testset "Complete Transform ChainRules" begin
Expand Down

0 comments on commit b16ba40

Please sign in to comment.