From 0578a7a68b397994ae5cd5f527730e2f77f6b785 Mon Sep 17 00:00:00 2001 From: rocroc2017 Date: Sat, 14 Dec 2024 18:23:12 +0800 Subject: [PATCH] Number format flexibility with set! (#634) * fix: a bug when using set! * update changelog * remove time stamp comments * assert number format before fourier transform * set_vordiv! with NF conversion for FFTW --------- Co-authored-by: Lu Peng <670123489@qq.com> Co-authored-by: milankl --- CHANGELOG.md | 1 + src/SpeedyTransforms/fourier.jl | 2 ++ src/SpeedyTransforms/spectral_transform.jl | 3 +++ src/dynamics/prognostic_variables.jl | 22 +++++++++++++++------- 4 files changed, 21 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9e2931dc4..300171602 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,7 @@ ## Unreleased +- Number format flexibility with set! [#634](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/634) - Forcing/drag for primitive models [#635](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/635) - RingGrids indexing with leading Colon should now always return another RingGrid instance [#637](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/637) - Roll back GPUArrays upgrade to ensure CUDA compatibility [#636](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/636) diff --git a/src/SpeedyTransforms/fourier.jl b/src/SpeedyTransforms/fourier.jl index 10bb60810..381b7ead5 100644 --- a/src/SpeedyTransforms/fourier.jl +++ b/src/SpeedyTransforms/fourier.jl @@ -33,6 +33,7 @@ function _fourier_batched!( # GRID TO SPECTRAL (; rfft_plans) = S # pre-planned transforms nlayers = size(grids, 2) # number of vertical layers + @assert eltype(grids) == eltype(S) "Number format of grid $(eltype(grids)) and SpectralTransform $(eltype(S)) need too match." @boundscheck ismatching(S, grids) || throw(DimensionMismatch(S, grids)) @boundscheck nlayers == S.nlayers || throw(DimensionMismatch(S, grids)) @boundscheck size(f_north) == size(f_south) == (S.nfreq_max, S.nlayers, nlat_half) || throw(DimensionMismatch(S, grids)) @@ -82,6 +83,7 @@ function _fourier_serial!( # GRID TO SPECTRAL rfft_plans = S.rfft_plans_1D # pre-planned transforms nlayers = size(grids, 2) # number of vertical layers + @assert eltype(grids) == eltype(S) "Number format of grid $(eltype(grids)) and SpectralTransform $(eltype(S)) need too match." @boundscheck ismatching(S, grids) || throw(DimensionMismatch(S, grids)) @boundscheck nlayers <= S.nlayers || throw(DimensionMismatch(S, grids)) @boundscheck size(f_north) == size(f_south) == (S.nfreq_max, S.nlayers, nlat_half) || throw(DimensionMismatch(S, grids)) diff --git a/src/SpeedyTransforms/spectral_transform.jl b/src/SpeedyTransforms/spectral_transform.jl index 82f11d671..9dce9d868 100644 --- a/src/SpeedyTransforms/spectral_transform.jl +++ b/src/SpeedyTransforms/spectral_transform.jl @@ -85,6 +85,9 @@ struct SpectralTransform{ eigenvalues⁻¹::VectorType # = -1/(l*(l+1)) end +# eltype of a transform is the number format used within +Base.eltype(S::SpectralTransform{NF}) where NF = NF + """ $(TYPEDSIGNATURES) Generator function for a SpectralTransform struct. With `NF` the number format, diff --git a/src/dynamics/prognostic_variables.jl b/src/dynamics/prognostic_variables.jl index e3dab304c..1197819f5 100644 --- a/src/dynamics/prognostic_variables.jl +++ b/src/dynamics/prognostic_variables.jl @@ -296,7 +296,13 @@ function set!( S::Union{Nothing, SpectralTransform}=nothing; add::Bool=false, ) - specs = isnothing(S) ? transform(grids) : transform(grids, S) + if isnothing(S) + specs = transform(grids) + else + # convert to number format in S, needed for FFTW + grids = convert.(eltype(S), grids) + specs = transform(grids, S) + end set!(var, specs; add) end @@ -440,15 +446,19 @@ function set_vordiv!( u::AbstractGridArray, v::AbstractGridArray, geometry::Geometry, - S::Union{Nothing, SpectralTransform}=nothing; + S::SpectralTransform = SpectralTransform(geometry.spectral_grid); add::Bool=false, coslat_scaling_included::Bool=false, ) u_ = coslat_scaling_included ? u : RingGrids.scale_coslat⁻¹(u) v_ = coslat_scaling_included ? v : RingGrids.scale_coslat⁻¹(v) - u_spec = isnothing(S) ? transform(u_) : transform(u_, S) - v_spec = isnothing(S) ? transform(v_) : transform(v_, S) + # convert to number format of spectral transform, otherwise FFTW complains + u_ = eltype(S) == eltype(u_) ? u_ : convert.(eltype(S), u_) + v_ = eltype(S) == eltype(v_) ? v_ : convert.(eltype(S), v_) + + u_spec = transform(u_, S) + v_spec = transform(v_, S) set_vordiv!(vor, div, u_spec, v_spec, geometry, S; add, coslat_scaling_included=true) end @@ -460,12 +470,10 @@ function set_vordiv!( u::LowerTriangularArray, v::LowerTriangularArray, geometry::Geometry, - S::Union{Nothing, SpectralTransform}=nothing; + S::SpectralTransform = SpectralTransform(geometry.spectral_grid); add::Bool=false, coslat_scaling_included::Bool=false, ) - S = isnothing(S) ? SpectralTransform(geometry.spectral_grid) : S - u_ = coslat_scaling_included ? u : transform(RingGrids.scale_coslat⁻¹(transform(u, S)), S) v_ = coslat_scaling_included ? v : transform(RingGrids.scale_coslat⁻¹(transform(u, S)), S)