Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Introduce separate extended tests in CI #628

Merged
merged 4 commits into from
Dec 4, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 29 additions & 0 deletions .github/workflows/extended_tests.yml
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
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

## Unreleased

- 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
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
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
Loading