Merge branch 'main' into mk/heating

CHANGELOG.md

@@ -3,6 +3,7 @@
 ## Unreleased
 
 - ConvectiveHeating implemented [#639](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/639)
+- 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)

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

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,

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)