Merge pull request #381 from SpeedyWeather/mk/output

scale_coslat funcs without Geometry, moved to RingGrids

src/RingGrids/RingGrids.jl

@@ -33,6 +33,7 @@ export  grids_match,
         get_nlat_half,
         get_npoints,
         get_latdlonds,
+        get_lat,
         get_colat,
         get_latd,
         get_lond,
@@ -46,6 +47,12 @@ export  grids_match,
         get_quadrature_weights,
         get_solid_angles
 
+# SCALING
+export  scale_coslat!,
+        scale_coslat²!,
+        scale_coslat⁻¹!,
+        scale_coslat⁻²!
+
 # INTERPOLATION
 export  AbstractInterpolator,
         GridGeometry,
@@ -72,4 +79,6 @@ include("quadrature_weights.jl")
 include("interpolation.jl")
 include("show.jl")
 include("similar.jl")
+include("scaling.jl")
+
 end

src/RingGrids/grids_general.jl

@@ -132,15 +132,16 @@ function get_latdlonds(Grid::Type{<:AbstractGrid},nlat_half::Integer)
     return latds, londs
 end
 
+get_lat(grid::Grid) where {Grid<:AbstractGrid} = get_lat(Grid,grid.nlat_half)
 get_latd(grid::Grid) where {Grid<:AbstractGrid} = get_latd(Grid,grid.nlat_half)
 get_lond(grid::Grid) where {Grid<:AbstractGrid} = get_lond(Grid,grid.nlat_half)
 
+function get_lat(Grid::Type{<:AbstractGrid},nlat_half::Integer)
+    return π/2 .- get_colat(Grid,nlat_half)
+end
+
 function get_latd(Grid::Type{<:AbstractGrid},nlat_half::Integer)
-    colat = get_colat(Grid,nlat_half)
-    latd = colat
-    latd .= π/2 .- colat
-    latd .= latd .* (360/2π)
-    return latd
+    return get_lat(Grid,nlat_half) * (360/2π)
 end
 
 # only defined for full grids, empty vector as fallback

src/RingGrids/scaling.jl

@@ -0,0 +1,29 @@
+# alias functions to scale the latitude of any gridded map A
+scale_coslat!(  A::AbstractGrid) = _scale_coslat!(A,power=1)
+scale_coslat²!( A::AbstractGrid) = _scale_coslat!(A,power=2)
+scale_coslat⁻¹!(A::AbstractGrid) = _scale_coslat!(A,power=-1)
+scale_coslat⁻²!(A::AbstractGrid) = _scale_coslat!(A,power=-2)
+
+function _scale_coslat!(A::Grid;power=1) where {Grid<:AbstractGrid}
+    coslat = sin.(get_colat(Grid,A.nlat_half))    # sin(colat) = cos(lat)
+    coslat .^= power
+    return _scale_lat!(A,coslat)
+end
+
+"""
+$(TYPEDSIGNATURES)
+Generic latitude scaling applied to `A` in-place with latitude-like vector `v`."""
+function _scale_lat!(A::AbstractGrid{NF},v::AbstractVector) where NF
+    @boundscheck get_nlat(A) == length(v) || throw(BoundsError)
+
+    rings = eachring(A)
+
+    @inbounds for (j,ring) in enumerate(rings)
+        vj = convert(NF,v[j])
+        for ij in ring
+            A[ij] *= vj
+        end
+    end
+
+    return A
+end 

src/SpeedyTransforms/spectral_transform.jl

@@ -501,11 +501,12 @@ on the truncation defined for `grid`. SpectralTransform struct `S` is allocated
 function gridded(   alms::AbstractMatrix{T};            # spectral coefficients
                     recompute_legendre::Bool = true,    # saves memory
                     Grid::Type{<:AbstractGrid} = DEFAULT_GRID,
+                    kwargs...
                     ) where {NF,T<:Complex{NF}}         # number format NF
 
     lmax, mmax = size(alms) .- 1                        # -1 for 0-based degree l, order m
     S = SpectralTransform(NF,Grid,lmax,mmax;recompute_legendre)
-    return gridded(alms,S)
+    return gridded(alms,S;kwargs...)
 end
 
 """
@@ -514,13 +515,14 @@ Spectral transform (spectral to grid space) from spherical coefficients `alms` t
 field `map` with precalculated properties based on the SpectralTransform struct `S`. `alms` is converted to
 a `LowerTriangularMatrix` to execute the in-place `gridded!`."""
 function gridded(   alms::AbstractMatrix,       # spectral coefficients
-                    S::SpectralTransform{NF},   # struct for spectral transform parameters
+                    S::SpectralTransform{NF};   # struct for spectral transform parameters
+                    kwargs...
                     ) where NF                  # number format NF
 
     map = zeros(S.Grid{NF},S.nlat_half)         # preallocate output
     almsᴸ = zeros(LowerTriangularMatrix{Complex{NF}},S.lmax+1,S.mmax+1)
     copyto!(almsᴸ,alms)                         # drop the upper triangle and convert to NF  
-    gridded!(map,almsᴸ,S)                       # now execute the in-place version
+    gridded!(map,almsᴸ,S;kwargs...)             # now execute the in-place version
     return map
 end
 

src/dynamics/initial_conditions.jl

@@ -484,7 +484,7 @@ function initialize_humidity!(  progn::PrognosticVariables,
     # @. humid_surf_grid = humid_ref*(exp(pres_surf_grid)/(pres_ref*100))^scale_height_ratio
     q_ref = 20e-3       # kg/kg at the surface
     @. humid_surf_grid .= q_ref
-    scale_coslat²!(humid_surf_grid,model.geometry)
+    RingGrids.scale_coslat²!(humid_surf_grid)
 
     humid_surf = spectral(humid_surf_grid,model.spectral_transform)
     spectral_truncation!(humid_surf)

src/dynamics/scaling.jl

@@ -1,31 +1,3 @@
-# alias functions to scale the latitude of any gridded map A
-scale_coslat!(  A::AbstractGrid,G::Geometry) = _scale_lat!(A,G.coslat)
-scale_coslat²!( A::AbstractGrid,G::Geometry) = _scale_lat!(A,G.coslat²)
-scale_coslat⁻¹!(A::AbstractGrid,G::Geometry) = _scale_lat!(A,G.coslat⁻¹)
-scale_coslat⁻²!(A::AbstractGrid,G::Geometry) = _scale_lat!(A,G.coslat⁻²)
-
-# matrix versions used for output
-scale_coslat!(  A::AbstractMatrix,G::Geometry) = A.*G.coslat'
-scale_coslat²!( A::AbstractMatrix,G::Geometry) = A.*G.coslat²'
-scale_coslat⁻¹!(A::AbstractMatrix,G::Geometry) = A.*G.coslat⁻¹'
-scale_coslat⁻²!(A::AbstractMatrix,G::Geometry) = A.*G.coslat⁻²'
-
-"""
-$(TYPEDSIGNATURES)
-Generic latitude scaling applied to `A` in-place with latitude-like vector `v`."""
-function _scale_lat!(A::AbstractGrid{NF},v::AbstractVector) where {NF<:AbstractFloat}
-    @boundscheck get_nlat(A) == length(v) || throw(BoundsError)
-
-    rings = eachring(A)
-
-    @inbounds for (j,ring) in enumerate(rings)
-        vj = convert(NF,v[j])
-        for ij in ring
-            A[ij] *= vj
-        end
-    end
-end 
-
 """
 $(TYPEDSIGNATURES)
 Scale the variable `var` inside `progn` with scalar `scale`.

test/netcdf_output.jl

@@ -5,55 +5,55 @@ import NetCDF
     n_days = 1
 
     # default grid, Float64, ShallowWater
-    spectral_grid = SpectralGrid(;NF=Float64)
+    spectral_grid = SpectralGrid(;NF=Float64,nlev=1)
     output = OutputWriter(spectral_grid,ShallowWater,path=tmp_output_path)
     model = ShallowWaterModel(;spectral_grid,output)
     simulation = initialize!(model)
     run!(simulation,output=true;n_days)
     @test simulation.model.feedback.nars_detected == false
 
     # default grid, Float32, ShallowWater
-    spectral_grid = SpectralGrid(;NF=Float32)
+    spectral_grid = SpectralGrid(;NF=Float32,nlev=1)
     output = OutputWriter(spectral_grid,ShallowWater,path=tmp_output_path)
     model = ShallowWaterModel(;spectral_grid,output)
     simulation = initialize!(model)
     run!(simulation,output=true;n_days)
     @test simulation.model.feedback.nars_detected == false
 
     # FullClenshawGrid, Float32, ShallowWater
-    spectral_grid = SpectralGrid(;NF=Float32,Grid=FullClenshawGrid)
+    spectral_grid = SpectralGrid(;NF=Float32,Grid=FullClenshawGrid,nlev=1)
     output = OutputWriter(spectral_grid,ShallowWater,path=tmp_output_path)
     model = ShallowWaterModel(;spectral_grid,output)
     simulation = initialize!(model)
     run!(simulation,output=true;n_days)
     @test simulation.model.feedback.nars_detected == false
 
     # OctahedralClenshawGrid, Float32, ShallowWater
-    spectral_grid = SpectralGrid(;NF=Float32,Grid=OctahedralClenshawGrid)
+    spectral_grid = SpectralGrid(;NF=Float32,Grid=OctahedralClenshawGrid,nlev=1)
     output = OutputWriter(spectral_grid,ShallowWater,path=tmp_output_path)
     model = ShallowWaterModel(;spectral_grid,output)
     simulation = initialize!(model)
     run!(simulation,output=true;n_days)
     @test simulation.model.feedback.nars_detected == false
 
     # HEALPixGrid, Float32, ShallowWater
-    spectral_grid = SpectralGrid(;NF=Float32,Grid=HEALPixGrid)
+    spectral_grid = SpectralGrid(;NF=Float32,Grid=HEALPixGrid,nlev=1)
     output = OutputWriter(spectral_grid,ShallowWater,path=tmp_output_path)
     model = ShallowWaterModel(;spectral_grid,output)
     simulation = initialize!(model)
     run!(simulation,output=true;n_days)
     @test simulation.model.feedback.nars_detected == false
 
     # OctaHEALPixGrid, Float32, ShallowWater
-    spectral_grid = SpectralGrid(;NF=Float32,Grid=OctaHEALPixGrid)
+    spectral_grid = SpectralGrid(;NF=Float32,Grid=OctaHEALPixGrid,nlev=1)
     output = OutputWriter(spectral_grid,ShallowWater,path=tmp_output_path)
     model = ShallowWaterModel(;spectral_grid,output)
     simulation = initialize!(model)
     run!(simulation,output=true;n_days)
     @test simulation.model.feedback.nars_detected == false
 
     # OctahedralClenshawGrid, as matrix, Float32, ShallowWater
-    spectral_grid = SpectralGrid(;NF=Float32,Grid=OctahedralClenshawGrid)
+    spectral_grid = SpectralGrid(;NF=Float32,Grid=OctahedralClenshawGrid,nlev=1)
     output = OutputWriter(spectral_grid,ShallowWater,path=tmp_output_path,as_matrix=true)
     model = ShallowWaterModel(;spectral_grid,output)
     simulation = initialize!(model)

test/run_speedy.jl

@@ -1,13 +1,13 @@
 @testset "run_speedy no errors, no blowup" begin
     # Barotropic
-    spectral_grid = SpectralGrid()
+    spectral_grid = SpectralGrid(nlev=1)
     model = BarotropicModel(;spectral_grid)
     simulation = initialize!(model)
     run!(simulation,n_days=10)
     @test simulation.model.feedback.nars_detected == false
 
     # ShallowWater
-    spectral_grid = SpectralGrid()
+    spectral_grid = SpectralGrid(nlev=1)
     model = ShallowWaterModel(;spectral_grid)
     simulation = initialize!(model)
     run!(simulation,n_days=10)

test/spectral_gradients.jl

@@ -27,8 +27,8 @@
 
         G = m.geometry
         S = m.spectral_transform
-        SpeedyWeather.scale_coslat⁻¹!(u_grid,G)
-        SpeedyWeather.scale_coslat⁻¹!(v_grid,G)
+        RingGrids.scale_coslat⁻¹!(u_grid)
+        RingGrids.scale_coslat⁻¹!(v_grid)
 
         uω_coslat⁻¹ = d.layers[1].dynamics_variables.a
         vω_coslat⁻¹ = d.layers[1].dynamics_variables.b
@@ -102,13 +102,13 @@ end
 
             A = Grid(randn(NF,SG.npoints))
             B = copy(A)
-            SpeedyWeather.scale_coslat⁻¹!(A,G)
-            SpeedyWeather.scale_coslat!(A,G)
+            RingGrids.scale_coslat⁻¹!(A)
+            RingGrids.scale_coslat!(A)
 
             @test all(isapprox.(A,B,rtol=10*eps(NF)))
 
-            SpeedyWeather.scale_coslat²!(A,G)
-            SpeedyWeather.scale_coslat⁻²!(A,G)
+            RingGrids.scale_coslat²!(A)
+            RingGrids.scale_coslat⁻²!(A)
 
             @test all(isapprox.(A,B,rtol=10*eps(NF)))
         end
@@ -219,8 +219,8 @@ end
 
         # times coslat² in grid space
         G = m.geometry
-        SpeedyWeather.scale_coslat⁻¹!(u,G)
-        SpeedyWeather.scale_coslat⁻¹!(v,G)
+        RingGrids.scale_coslat⁻¹!(u)
+        RingGrids.scale_coslat⁻¹!(v)
 
         # transform back
         S = m.spectral_transform
@@ -310,8 +310,8 @@ end
         dadx_grid = gridded(dadx,m.spectral_transform)
         dady_grid = gridded(dady,m.spectral_transform)
 
-        SpeedyWeather.scale_coslat⁻²!(dadx_grid,m.geometry)
-        SpeedyWeather.scale_coslat⁻²!(dady_grid,m.geometry)
+        RingGrids.scale_coslat⁻²!(dadx_grid)
+        RingGrids.scale_coslat⁻²!(dady_grid)
 
         SpeedyWeather.spectral!(dadx,dadx_grid,m.spectral_transform)
         SpeedyWeather.spectral!(dady,dady_grid,m.spectral_transform)