diff --git a/Project.toml b/Project.toml index 1598f4975..2ad23b26f 100644 --- a/Project.toml +++ b/Project.toml @@ -31,9 +31,11 @@ UnicodePlots = "b8865327-cd53-5732-bb35-84acbb429228" [weakdeps] Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a" +JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb" [extensions] SpeedyWeatherMakieExt = "Makie" +SpeedyWeatherJLArraysExt = "JLArrays" [compat] AbstractFFTs = "1" @@ -47,7 +49,9 @@ DocStringExtensions = "0.9" FFTW = "1" FLoops = "0.2" FastGaussQuadrature = "0.4, 0.5, 1" +GPUArrays = "10" GenericFFT = "0.1" +JLArrays = "0.1.4" JLD2 = "0.4" KernelAbstractions = "0.9" LinearAlgebra = "1.9" @@ -63,8 +67,8 @@ UnicodePlots = "3.3" julia = "1.9" [extras] -Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] test = ["Test", "JLArrays"] diff --git a/docs/src/output.md b/docs/src/output.md index 6c72287a4..56ff7a29e 100644 --- a/docs/src/output.md +++ b/docs/src/output.md @@ -96,13 +96,13 @@ for the model integration. ```julia my_output_writer = OutputWriter(spectral_grid, ShallowWater, output_Grid=FullClenshawGrid, nlat_half=48) ``` -Note that by default the output is on the corresponding full of the grid used in the dynamical core +Note that by default the output is on the corresponding full type of the grid type used in the dynamical core so that interpolation only happens at most in the zonal direction as they share the location of the latitude rings. You can check this by ```@example netcdf -RingGrids.full_grid(OctahedralGaussianGrid) +RingGrids.full_grid_type(OctahedralGaussianGrid) ``` -So the corresponding full grid of an `OctahedralGaussianGrid` is the `FullGaussiangrid` and the same resolution +So the corresponding full grid of an `OctahedralGaussianGrid` is the `FullGaussianGrid` and the same resolution `nlat_half` is chosen by default in the output writer (which you can change though as shown above). Overview of the corresponding full grids diff --git a/ext/SpeedyWeatherJLArraysExt.jl b/ext/SpeedyWeatherJLArraysExt.jl new file mode 100644 index 000000000..8f33c861e --- /dev/null +++ b/ext/SpeedyWeatherJLArraysExt.jl @@ -0,0 +1,10 @@ +module SpeedyWeatherJLArraysExt + +using SpeedyWeather, JLArrays + +# for RingGrids and LowerTriangularMatrices: +# every Array needs this method to strip away the parameters +SpeedyWeather.RingGrids.nonparametric_type(::Type{<:JLArray}) = JLArray +SpeedyWeather.LowerTriangularMatrices.nonparametric_type(::Type{<:JLArray}) = JLArray + +end # module \ No newline at end of file diff --git a/ext/SpeedyWeatherMakieExt.jl b/ext/SpeedyWeatherMakieExt.jl index fd09dd8e4..d355318a5 100644 --- a/ext/SpeedyWeatherMakieExt.jl +++ b/ext/SpeedyWeatherMakieExt.jl @@ -12,7 +12,7 @@ function Makie.heatmap( title::String = "$(RingGrids.get_nlat(grid))-ring $(typeof(grid))", kwargs... # pass on to Makie.heatmap ) - full_grid = RingGrids.interpolate(RingGrids.full_grid(typeof(grid)), grid.nlat_half, grid) + full_grid = RingGrids.interpolate(RingGrids.full_grid_type(grid), grid.nlat_half, grid) heatmap(full_grid; title, kwargs...) end diff --git a/src/RingGrids/RingGrids.jl b/src/RingGrids/RingGrids.jl index 4dc5b857b..a5d46952a 100644 --- a/src/RingGrids/RingGrids.jl +++ b/src/RingGrids/RingGrids.jl @@ -9,50 +9,77 @@ import Statistics: mean import FastGaussQuadrature import LinearAlgebra -# GRIDS +# GPU +import Adapt +import GPUArrays +import CUDA + +# ABSTRACT GRIDS (2D) AND GRIDARRAYS (3D+) +export AbstractGridArray, + AbstractFullGridArray, + AbstractReducedGridArray + export AbstractGrid, AbstractFullGrid, - AbstractOctahedralGrid, - AbstractHEALPixGrid, - AbstractOctaHEALPixGrid + AbstractReducedGrid + +# CONCRETE GRIDS (2D) AND GRIDARRAYS (3D+) +export FullGaussianArray, + FullClenshawArray, + FullHEALPixArray, + FullOctaHEALPixArray export FullGaussianGrid, FullClenshawGrid, FullHEALPixGrid, FullOctaHEALPixGrid +export OctahedralGaussianArray, + OctahedralClenshawArray, + HEALPixArray, + OctaHEALPixArray + export OctahedralGaussianGrid, OctahedralClenshawGrid, HEALPixGrid, OctaHEALPixGrid -# GRID FUNCTIONS +# SIZE export grids_match, - get_truncation, - get_resolution, get_nlat, get_nlat_half, - get_npoints, - get_latdlonds, + get_npoints2D + +# COORDINATES +export get_latdlonds, get_lat, get_colat, get_latd, get_lond, - each_index_in_ring, - each_index_in_ring!, - eachgridpoint, - eachring, - whichring, get_nlons, - get_nlon_max, - get_quadrature_weights, + get_nlon_max + +# INTEGRATION +export get_quadrature_weights, get_solid_angles +# ITERATORS +export eachgrid, + eachring, + whichring, + eachgridpoint, + each_index_in_ring, + each_index_in_ring! + # SCALING export scale_coslat!, scale_coslat²!, scale_coslat⁻¹!, - scale_coslat⁻²! + scale_coslat⁻²!, + scale_coslat, + scale_coslat², + scale_coslat⁻¹, + scale_coslat⁻² # INTERPOLATION export AbstractInterpolator, @@ -67,19 +94,31 @@ export interpolate, update_locator, update_locator! -# export plot - include("utility_functions.jl") -include("grids_general.jl") +# GENERAL +include("general.jl") include("full_grids.jl") -include("octahedral.jl") -include("healpix.jl") -include("octahealpix.jl") +include("reduced_grids.jl") +include("scaling.jl") + +# FULL GRIDS +include("grids/full_gaussian.jl") +include("grids/full_clenshaw.jl") +include("grids/full_healpix.jl") +include("grids/full_octahealpix.jl") + +# REDUCED GRIDS +include("grids/octahedral_gaussian.jl") +include("grids/octahedral_clenshaw.jl") +include("grids/healpix.jl") +include("grids/octahealpix.jl") + +# INTEGRATION AND INTERPOLATION include("quadrature_weights.jl") include("interpolation.jl") + +# OUTPUT include("show.jl") -include("similar.jl") -include("scaling.jl") end \ No newline at end of file diff --git a/src/RingGrids/full_grids.jl b/src/RingGrids/full_grids.jl index 84ab8d904..2b9404305 100644 --- a/src/RingGrids/full_grids.jl +++ b/src/RingGrids/full_grids.jl @@ -1,42 +1,77 @@ -""" - abstract type AbstractFullGrid{T} <: AbstractGrid{T} end +"""Subtype of `AbstractGridArray` for all N-dimensional arrays of ring grids that have the +same number of longitude points on every ring. As such these (horizontal) grids are representable +as a matrix, with denser grid points towards the poles.""" +abstract type AbstractFullGridArray{T, N, ArrayType <: AbstractArray{T, N}} <: AbstractGridArray{T, N, ArrayType} end -An `AbstractFullGrid` is a horizontal grid with a constant number of longitude +"""An `AbstractFullGrid` is a horizontal grid with a constant number of longitude points across latitude rings. Different latitudes can be used, Gaussian latitudes, -equi-angle latitdes, or others.""" -abstract type AbstractFullGrid{T} <: AbstractGrid{T} end +equi-angle latitudes (also called Clenshaw from Clenshaw-Curtis quadrature), or others.""" +const AbstractFullGrid{T} = AbstractFullGridArray{T, 1, Vector{T}} +full_grid_type(Grid::Type{<:AbstractFullGridArray}) = horizontal_grid_type(Grid) +full_array_type(Grid::Type{<:AbstractFullGridArray}) = nonparametric_type(Grid) + +## SIZE +get_nlon_max(Grid::Type{<:AbstractFullGridArray}, nlat_half::Integer) = get_nlon(Grid, nlat_half) +get_nlon_per_ring(Grid::Type{<:AbstractFullGridArray}, nlat_half::Integer, j::Integer) = + get_nlon(Grid, nlat_half) +matrix_size(Grid::Type{<:AbstractFullGridArray}, nlat_half::Integer) = + (get_nlon(Grid, nlat_half), get_nlat(Grid, nlat_half)) + +## CONVERSION +# convert an AbstractMatrix to the full grids, and vice versa +(Grid::Type{<:AbstractFullGrid})(M::AbstractMatrix) = Grid(vec(M)) +Base.Array(grid::AbstractFullGridArray) = Array(reshape(grid.data, :, get_nlat(grid), size(grid.data)[2:end]...)) +Base.Matrix(grid::AbstractFullGridArray) = Array(grid) + +## INDEXING + +"""$(TYPEDSIGNATURES) `UnitRange` for every grid point of grid `Grid` of resolution `nlat_half` +on ring `j` (`j=1` is closest ring around north pole, `j=nlat` around south pole).""" +function each_index_in_ring( + Grid::Type{<:AbstractFullGridArray}, # function for full grids + j::Integer, # ring index north to south + nlat_half::Integer, +) + @boundscheck 0 < j <= get_nlat(Grid, nlat_half) || throw(BoundsError) # valid ring index? + nlon = get_nlon(Grid, nlat_half) # number of longitudes per ring (const) + index_1st = (j-1)*nlon + 1 # first in-ring index i + index_end = j*nlon # last in-ring index i + return index_1st:index_end # range of js in ring +end -get_nlon(Grid::Type{<:AbstractFullGrid}, nlat_half::Integer) = get_nlon_max(Grid, nlat_half) -get_nlon_max(::Type{<:AbstractFullGrid}, nlat_half::Integer) = 4nlat_half -get_nlon_per_ring(Grid::Type{<:AbstractFullGrid}, nlat_half::Integer, j::Integer) = - get_nlon_max(Grid, nlat_half) +# precompute ring indices for full grids +function each_index_in_ring!( + rings::Vector{<:UnitRange{<:Integer}}, + Grid::Type{<:AbstractFullGridArray}, + nlat_half::Integer, +) + nlat = length(rings) # number of latitude rings + @boundscheck nlat == get_nlat(Grid, nlat_half) || throw(BoundsError) -function get_lon(Grid::Type{<:AbstractFullGrid}, nlat_half::Integer) - nlat_half == 0 && return Float64[] - nlon = get_nlon(Grid, nlat_half) - return collect(range(0, 2π-π/nlon, step=2π/nlon)) + nlon = get_nlon(Grid, nlat_half) # number of longitudes + index_end = 0 + @inbounds for j in 1:nlat + index_1st = index_end + 1 # 1st index is +1 from prev ring's last index + index_end += nlon # only calculate last index per ring + rings[j] = index_1st:index_end # write UnitRange to rings vector + end end -function get_lond(Grid::Type{<:AbstractFullGrid}, nlat_half::Integer) +## COORDINATES +function get_lond(Grid::Type{<:AbstractFullGridArray}, nlat_half::Integer) lon = get_lon(Grid, nlat_half) lon .*= 360/2π # convert to lond in-place return lon # = lond end -# convert an AbstractMatrix to the full grids, and vice versa -(Grid::Type{<:AbstractFullGrid})(M::AbstractMatrix{T}) where T = Grid{T}(vec(M)) -Base.Matrix(grid::AbstractFullGrid{T}) where T = Matrix{T}(reshape(grid.data, :, get_nlat(grid))) -matrix_size(grid::AbstractFullGrid) = (get_nlon_max(grid), get_nlat(grid)) -matrix_size(Grid::Type{<:AbstractFullGrid}, n::Integer) = (get_nlon_max(Grid, n), get_nlat(Grid, n)) - -function get_colatlons(Grid::Type{<:AbstractFullGrid}, nlat_half::Integer) +function get_colatlons(Grid::Type{<:AbstractFullGridArray}, nlat_half::Integer) - colat = get_colat(Grid, nlat_half) # vector of colats [0, π] - lon = get_lon(Grid, nlat_half) # vector of longitudes [0, 2π) - nlon = get_nlon(Grid, nlat_half) # number of longitudes - nlat = get_nlat(Grid, nlat_half) # number of latitudes + colat = get_colat(Grid, nlat_half) # vector of colats [0, π] + lon = get_lon(Grid, nlat_half) # vector of longitudes [0, 2π) + nlon = get_nlon(Grid, nlat_half) # number of longitudes + nlat = get_nlat(Grid, nlat_half) # number of latitudes - npoints = get_npoints(Grid, nlat_half) # total number of grid points + npoints = get_npoints(Grid, nlat_half) # total number of grid points colats = zeros(npoints) # preallocate lons = zeros(npoints) @@ -49,156 +84,4 @@ function get_colatlons(Grid::Type{<:AbstractFullGrid}, nlat_half::Integer) end return colats, lons -end - -function each_index_in_ring(Grid::Type{<:AbstractFullGrid}, # function for full grids - j::Integer, # ring index north to south - nlat_half::Integer) # resolution param - - @boundscheck 0 < j <= get_nlat(Grid, nlat_half) || throw(BoundsError) # valid ring index? - nlon = 4nlat_half # number of longitudes per ring (const) - index_1st = (j-1)*nlon + 1 # first in-ring index i - index_end = j*nlon # last in-ring index i - return index_1st:index_end # range of js in ring -end - -function each_index_in_ring!( rings::Vector{<:UnitRange{<:Integer}}, - Grid::Type{<:AbstractFullGrid}, - nlat_half::Integer) - nlat = length(rings) # number of latitude rings - @boundscheck nlat == get_nlat(Grid, nlat_half) || throw(BoundsError) - - nlon = get_nlon(Grid, nlat_half) # number of longitudes - index_end = 0 - @inbounds for j in 1:nlat - index_1st = index_end + 1 # 1st index is +1 from prev ring's last index - index_end += nlon # only calculate last index per ring - rings[j] = index_1st:index_end # write UnitRange to rings vector - end -end - - -""" - G = FullClenshawGrid{T} - -A FullClenshawGrid is a regular latitude-longitude grid with an odd number of `nlat` equi-spaced -latitudes, the central latitude ring is on the Equator. The same `nlon` longitudes for every latitude ring. -The grid points are closer in zonal direction around the poles. The values of all grid points are stored -in a vector field `data` that unravels the data 0 to 360˚, then ring by ring, which are sorted north to south.""" -struct FullClenshawGrid{T} <: AbstractFullGrid{T} - data::Vector{T} # data vector, ring by ring, north to south - nlat_half::Int # number of latitudes on one hemisphere (incl Equator) - - FullClenshawGrid{T}(data::AbstractVector, nlat_half::Integer) where T = length(data) == npoints_clenshaw(nlat_half) ? - new(data, nlat_half) : error("$(length(data))-element Vector{$(eltype(data))} cannot be used to create a "* - "L$nlat_half ($(4nlat_half)x$(2nlat_half - 1)) FullClenshawGrid{$T}.") -end - -nonparametric_type(::Type{<:FullClenshawGrid}) = FullClenshawGrid - -# subtract the otherwise double-counted 4nlat_half equator points -npoints_clenshaw(nlat_half::Integer) = 8nlat_half^2 - 4nlat_half -nlat_half_clenshaw(npoints::Integer) = round(Int, 1/4 + sqrt(1/16 + npoints/8)) # inverse - -# infer nlat_half from data vector length, infer parametric type from eltype of data -FullClenshawGrid{T}(data::AbstractVector) where T = FullClenshawGrid{T}(data, nlat_half_clenshaw(length(data))) -FullClenshawGrid(data::AbstractVector, n::Integer...) = FullClenshawGrid{eltype(data)}(data, n...) - -nlat_odd(::Type{<:FullClenshawGrid}) = true -get_npoints(::Type{<:FullClenshawGrid}, nlat_half::Integer) = npoints_clenshaw(nlat_half) -get_colat(::Type{<:FullClenshawGrid}, nlat_half::Integer) = [j/(2nlat_half)*π for j in 1:2nlat_half-1] -get_quadrature_weights(::Type{<:FullClenshawGrid}, nlat_half::Integer) = clenshaw_curtis_weights(nlat_half) -full_grid(::Type{<:FullClenshawGrid}) = FullClenshawGrid # the full grid with same latitudes - -""" - G = FullGaussianGrid{T} - -A full Gaussian grid is a regular latitude-longitude grid that uses `nlat` Gaussian latitudes, -and the same `nlon` longitudes for every latitude ring. The grid points are closer in zonal direction -around the poles. The values of all grid points are stored in a vector field `v` that unravels -the data 0 to 360˚, then ring by ring, which are sorted north to south.""" -struct FullGaussianGrid{T} <: AbstractFullGrid{T} - data::Vector{T} # data vector, ring by ring, north to south - nlat_half::Int # number of latitudes on one hemisphere - - FullGaussianGrid{T}(data::AbstractVector, nlat_half::Integer) where T = length(data) == 8nlat_half^2 ? - new(data, nlat_half) : error("$(length(data))-element Vector{$(eltype(data))} cannot be used to create a "* - "F$nlat_half ($(4nlat_half)x$(2nlat_half)) FullGaussianGrid{$T}.") -end - -nonparametric_type(::Type{<:FullGaussianGrid}) = FullGaussianGrid - -npoints_gaussian(nlat_half::Integer) = 8nlat_half^2 -nlat_half_gaussian(npoints::Integer) = round(Int, sqrt(npoints/8)) - -# infer nlat_half from data vector length, infer parametric type from eltype of data -FullGaussianGrid{T}(data::AbstractVector) where T = FullGaussianGrid{T}(data, nlat_half_gaussian(length(data))) -FullGaussianGrid(data::AbstractVector, n::Integer...) = FullGaussianGrid{eltype(data)}(data, n...) - -nlat_odd(::Type{<:FullGaussianGrid}) = false -get_npoints(::Type{<:FullGaussianGrid}, nlat_half::Integer) = npoints_gaussian(nlat_half) -get_colat(::Type{<:FullGaussianGrid}, nlat_half::Integer) = - π .- acos.(FastGaussQuadrature.gausslegendre(2nlat_half)[1]) -get_quadrature_weights(::Type{<:FullGaussianGrid}, nlat_half::Integer) = gaussian_weights(nlat_half) -full_grid(::Type{<:FullGaussianGrid}) = FullGaussianGrid # the full grid with same latitudes - -""" - G = FullHEALPixGrid{T} - -A full HEALPix grid is a regular latitude-longitude grid that uses `nlat` latitudes from the HEALPix grid, -and the same `nlon` longitudes for every latitude ring. The grid points are closer in zonal direction -around the poles. The values of all grid points are stored in a vector field `v` that unravels -the data 0 to 360˚, then ring by ring, which are sorted north to south.""" -struct FullHEALPixGrid{T} <: AbstractFullGrid{T} - data::Vector{T} # data vector, ring by ring, north to south - nlat_half::Int # number of latitudes on one hemisphere - - FullHEALPixGrid{T}(data::AbstractVector, nlat_half::Integer) where T = length(data) == npoints_fullhealpix(nlat_half) ? - new(data, nlat_half) : error("$(length(data))-element Vector{$(eltype(data))} cannot be used to create a "* - "H$nlat_half ($(4nlat_half)x$(2nlat_half-1)) FullHEALPixGrid{$T}.") -end - -nonparametric_type(::Type{<:FullHEALPixGrid}) = FullHEALPixGrid - -npoints_fullhealpix(nlat_half::Integer) = 4nlat_half*(2nlat_half-1) -nlat_half_fullhealpix(npoints::Integer) = round(Int, 1/4 + sqrt(1/16 + npoints/8)) - -# infer nlat_half from data vector length, infer parametric type from eltype of data -FullHEALPixGrid{T}(data::AbstractVector) where T = FullHEALPixGrid{T}(data, nlat_half_fullhealpix(length(data))) -FullHEALPixGrid(data::AbstractVector, n::Integer...) = FullHEALPixGrid{eltype(data)}(data, n...) -nlat_odd(::Type{<:FullHEALPixGrid}) = true -get_npoints(::Type{<:FullHEALPixGrid}, nlat_half::Integer) = npoints_fullhealpix(nlat_half) -get_colat(::Type{<:FullHEALPixGrid}, nlat_half::Integer) = get_colat(HEALPixGrid, nlat_half) -full_grid(::Type{<:FullHEALPixGrid}) = FullHEALPixGrid # the full grid with same latitudes -get_quadrature_weights(::Type{<:FullHEALPixGrid}, nlat_half::Integer) = healpix_weights(nlat_half) - -""" - G = FullOctaHEALPixGrid{T} - -A full OctaHEALPix grid is a regular latitude-longitude grid that uses `nlat` OctaHEALPix latitudes, -and the same `nlon` longitudes for every latitude ring. The grid points are closer in zonal direction -around the poles. The values of all grid points are stored in a vector field `v` that unravels -the data 0 to 360˚, then ring by ring, which are sorted north to south.""" -struct FullOctaHEALPixGrid{T} <: AbstractFullGrid{T} - data::Vector{T} # data vector, ring by ring, north to south - nlat_half::Int # number of latitudes on one hemisphere - - FullOctaHEALPixGrid{T}(data::AbstractVector, nlat_half::Integer) where T = length(data) == npoints_fulloctahealpix(nlat_half) ? - new(data, nlat_half) : error("$(length(data))-element Vector{$(eltype(data))} cannot be used to create a "* - "F$nlat_half ($(4nlat_half)x$(2nlat_half - 1)) FullOctaHEALPixGrid{$T}.") -end - -nonparametric_type(::Type{<:FullOctaHEALPixGrid}) = FullOctaHEALPixGrid - -npoints_fulloctahealpix(nlat_half::Integer) = 8nlat_half^2 - 4nlat_half -nlat_half_fulloctahealpix(npoints::Integer) = round(Int, 1/4 + sqrt(1/16 + npoints/8)) - -# infer nlat_half from data vector length, infer parametric type from eltype of data -FullOctaHEALPixGrid{T}(data::AbstractVector) where T = FullOctaHEALPixGrid{T}(data, nlat_half_fulloctahealpix(length(data))) -FullOctaHEALPixGrid(data::AbstractVector, n::Integer...) = FullOctaHEALPixGrid{eltype(data)}(data, n...) - -nlat_odd(::Type{<:FullOctaHEALPixGrid}) = true -get_npoints(::Type{<:FullOctaHEALPixGrid}, nlat_half::Integer) = npoints_fulloctahealpix(nlat_half) -get_colat(::Type{<:FullOctaHEALPixGrid}, nlat_half::Integer) = get_colat(OctaHEALPixGrid, nlat_half) -get_quadrature_weights(::Type{<:FullOctaHEALPixGrid}, nlat_half::Integer) = octahealpix_weights(nlat_half) -full_grid(::Type{<:FullOctaHEALPixGrid}) = FullOctaHEALPixGrid # the full grid with same latitudes \ No newline at end of file +end \ No newline at end of file diff --git a/src/RingGrids/general.jl b/src/RingGrids/general.jl new file mode 100644 index 000000000..b26515382 --- /dev/null +++ b/src/RingGrids/general.jl @@ -0,0 +1,422 @@ +"""Abstract supertype for all arrays of ring grids, representing `N`-dimensional +data on the sphere in two dimensions (but unravelled into a vector in the first dimension, +the actual "ring grid") plus additional `N-1` dimensions for the vertical and/or time etc. +Parameter `T` is the `eltype` of the underlying data, held as in the array type `ArrayType` +(Julia's `Array` for CPU or others for GPU). + +Ring grids have several consecuitive grid points share the same latitude (= a ring), +grid points on a given ring are equidistant. Grid points are ordered 0 to 360˚E, +starting around the north pole, ring by ring to the south pole. """ +abstract type AbstractGridArray{T, N, ArrayType <: AbstractArray{T, N}} <: AbstractArray{T, N} end + +"""Abstract supertype for all ring grids, representing 2-dimensional data on the +sphere unravelled into a Julia `Vector`. Subtype of `AbstractGridArray` with +`N=1` and `ArrayType=Vector{T}` of `eltype T`.""" +const AbstractGrid{T} = AbstractGridArray{T, 1, Vector{T}} + +## TYPES +"""$(TYPEDSIGNATURES) For any instance of `AbstractGridArray` type its n-dimensional type +(*Grid{T, N, ...} returns *Array) but without any parameters `{T, N, ArrayType}`""" +nonparametric_type(grid::AbstractGridArray) = nonparametric_type(typeof(grid)) + +# also needed for other array types +nonparametric_type(::Type{Array}) = Array +nonparametric_type(::Type{CUDA.CuArray}) = CuArray + +"""$(TYPEDSIGNATURES) Full grid array type for `grid`. Always returns the N-dimensional `*Array` +not the two-dimensional (`N=1`) `*Grid`. For reduced grids the corresponding full grid that +share the same latitudes.""" +full_array_type(grid::AbstractGridArray) = full_array_type(typeof(grid)) + +"""$(TYPEDSIGNATURES) Full (horizontal) grid type for `grid`. Always returns the two-dimensional +(`N=1`) `*Grid` type. For reduced grids the corresponding full grid that share the same latitudes.""" +full_grid_type(grid::AbstractGridArray) = horizontal_grid_type(full_array_type(grid)) +full_grid_type(Grid::Type{<:AbstractGridArray}) = horizontal_grid_type(full_array_type(Grid)) + +"""$(TYPEDSIGNATURES) The two-dimensional (`N=1`) `*Grid` for `grid`, which can be an N-dimensional +`*GridArray`.""" +horizontal_grid_type(grid::AbstractGridArray) = horizontal_grid_type(typeof(grid)) + +## SIZE +Base.length(G::AbstractGridArray) = length(G.data) # total number of grid points +Base.size(G::AbstractGridArray) = size(G.data) # size of underlying data (horizontal is unravelled) + +"""$(TYPEDSIGNATURES) Size of underlying data array plus precomputed ring indices.""" +Base.sizeof(G::AbstractGridArray) = sizeof(G.data) + sizeof(G.rings) + +"""$(TYPEDSIGNATURES) Resolution paraemeters `nlat_half` of a `grid`. +Number of latitude rings on one hemisphere, Equator included.""" +get_nlat_half(grid::AbstractGridArray) = grid.nlat_half + +"""$(TYPEDSIGNATURES) True for a `grid` that has an odd number of +latitude rings `nlat` (both hemispheres).""" +nlat_odd(grid::AbstractGridArray) = nlat_odd(typeof(grid)) + +# get total number of latitude rings, *(nlat_half > 0) to return 0 for nlat_half = 0 +get_nlat(Grid::Type{<:AbstractGridArray}, nlat_half::Integer) = 2nlat_half - nlat_odd(Grid)*(nlat_half > 0) + +"""$(TYPEDSIGNATURES) Get number of latitude rings, pole to pole.""" +get_nlat(grid::Grid) where {Grid<:AbstractGridArray} = get_nlat(Grid, grid.nlat_half) + +"""$(TYPEDSIGNATURES) Total number of grid points in all dimensions of `grid`. +Equivalent to length of the underlying data array.""" +get_npoints(grid::Grid) where {Grid<:AbstractGridArray} = get_npoints(Grid, grid.nlat_half) +get_npoints(G::Type{<:AbstractGridArray}, nlat_half::Integer, k::Integer...) = prod(k) * get_npoints2D(G, nlat_half) + +"""$(TYPEDSIGNATURES) Number of grid points in the horizontal dimension only, +even if `grid` is N-dimensional.""" +get_npoints2D(grid::Grid) where {Grid<:AbstractGridArray} = get_npoints2D(Grid, grid.nlat_half) + +"""$(TYPEDSIGNATURES) Size of the matrix of the horizontal grid if representable as such (not all grids).""" +matrix_size(grid::Grid) where {Grid<:AbstractGridArray} = matrix_size(Grid, grid.nlat_half) + +## INDEXING +# simply propagate all indices forward +Base.@propagate_inbounds Base.getindex(G::AbstractGridArray, ijk...) = getindex(G.data, ijk...) + +@inline function Base.getindex( + G::GridArray, + col::Colon, + k::Integer..., +) where {GridArray<:AbstractGridArray} + GridArray_ = nonparametric_type(GridArray) # obtain parameters from G.data + return GridArray_(getindex(G.data, col, k...), G.nlat_half, G.rings) +end + +# simply propagate all indices forward +Base.@propagate_inbounds Base.setindex!(G::AbstractGridArray, x, ijk...) = setindex!(G.data, x, ijk...) +Base.fill!(G::AbstractGridArray, x) = fill!(G.data, x) + +## CONSTRUCTORS +"""$(TYPEDSIGNATURES) True for `data`, `nlat_half` and `rings` that all match in size +to construct a grid of type `Grid`.""" +function check_inputs(data, nlat_half, rings, Grid) + check = true + check &= size(data, 1) == get_npoints2D(Grid, nlat_half) # test number of 2D grid points + check &= length(rings) == get_nlat(Grid, nlat_half) # test number of rings == nlat + # TODO also check that rings map to all and only valid grid points? + return check +end + +function error_message(data, nlat_half, rings, G, T, N, A) + nlat = get_nlat(G, nlat_half) + nrings = length(rings) + if nlat != nrings + return error("$nrings-element ring indices "* + "cannot be used to create a $nlat-ring $G{$T, $N, $A}.") + else + return error("$(summary(data)) cannot be used to create a $nlat-ring $G{$T, $N, $A}") + end +end + +# if no rings are provided, calculate them +function (::Type{Grid})(data::AbstractArray, nlat_half::Integer) where {Grid<:AbstractGridArray} + GridArray_ = nonparametric_type(Grid) # strip away parameters of type, obtain parameters from data + rings = eachring(Grid, nlat_half) # precompute indices to access the variable-length rings + return GridArray_(data, nlat_half, rings) +end + +# if no nlat_half provided calculate it +function (::Type{Grid})(data::AbstractArray) where {Grid<:AbstractGridArray} + npoints2D = size(data, 1) # from 1st dim of data + nlat_half = get_nlat_half(Grid, npoints2D) # get nlat_half of Grid + return Grid(data, nlat_half) +end + +for f in (:zeros, :ones, :rand, :randn) + @eval begin + # general version with ArrayType(zeros(...)) conversion + function Base.$f( + ::Type{Grid}, + nlat_half::Integer, + k::Integer..., + ) where {Grid<:AbstractGridArray{T, N, ArrayType}} where {T, N, ArrayType} + return Grid(ArrayType($f(T, get_npoints2D(Grid, nlat_half), k...)), nlat_half) + end + + # CPU version with zeros(T, ...) producing Array + function Base.$f( + ::Type{Grid}, + nlat_half::Integer, + k::Integer..., + ) where {Grid<:AbstractGridArray{T}} where T + return Grid($f(T, get_npoints2D(Grid, nlat_half), k...), nlat_half) + end + + # use Float64 if no type provided + function Base.$f( + ::Type{Grid}, + nlat_half::Integer, + k::Integer..., + ) where {Grid<:AbstractGridArray} + return $f(Grid{Float64}, nlat_half, k...) + end + end +end + +# zero element of an AbstractGridArray instance grid by creating new zero(grid.data) +Base.zero(grid::Grid) where {Grid<:AbstractGridArray} = + nonparametric_type(Grid)(zero(grid.data), grid.nlat_half, grid.rings) + +# similar data but everything else identical +function Base.similar(grid::Grid) where {Grid<:AbstractGridArray} + return nonparametric_type(Grid)(similar(grid.data), grid.nlat_half, grid.rings) +end + +# data with new type T but everything else identical +function Base.similar(grid::Grid, ::Type{T}) where {Grid<:AbstractGridArray, T} + return nonparametric_type(Grid)(similar(grid.data, T), grid.nlat_half, grid.rings) +end + +# data with same type T but new size +function Base.similar( + grid::Grid, + nlat_half::Integer, + k::Integer... +) where {Grid<:AbstractGridArray{T, N, ArrayType}} where {T, N, ArrayType} + similar_data = similar(grid.data, get_npoints2D(Grid, nlat_half), k...) + return nonparametric_type(Grid)(similar_data, nlat_half) +end + +# data with new type T and new size +function Base.similar( + grid::Grid, + ::Type{Tnew}, + nlat_half::Integer, + k::Integer... +) where {Grid<:AbstractGridArray, Tnew} + similar_data = similar(grid.data, Tnew, get_npoints2D(Grid, nlat_half), k...) + return nonparametric_type(Grid)(similar_data, nlat_half) +end + +# general version with ArrayType{T, N}(undef, ...) generator +function (::Type{Grid})( + ::UndefInitializer, + nlat_half::Integer, + k::Integer..., +) where {Grid<:AbstractGridArray{T, N, ArrayType}} where {T, N, ArrayType} + ArrayType_ = nonparametric_type(ArrayType) + return Grid(ArrayType_{T, N}(undef, get_npoints2D(Grid, nlat_half), k...), nlat_half) +end + +# CPU version with Array{T, N}(undef, ...) generator +function (::Type{Grid})( + ::UndefInitializer, + nlat_half::Integer, + k::Integer..., +) where {Grid<:AbstractGridArray{T}} where T + return Grid(Array{T}(undef, get_npoints2D(Grid, nlat_half), k...), nlat_half) +end + +# use Float64 if no type provided +function (::Type{Grid})( + ::UndefInitializer, + nlat_half::Integer, + k::Integer..., +) where {Grid<:AbstractGridArray} + return Grid(Array{Float64}(undef, get_npoints2D(Grid, nlat_half), k...), nlat_half) +end + +## COORDINATES + +"""$(TYPEDSIGNATURES) Latitudes (in degrees, -90˚-90˚N) and longitudes (0-360˚E) for +every (horizontal) grid point in `grid`. Ordered 0-360˚E then north to south.""" +get_latdlonds(grid::Grid) where {Grid<:AbstractGridArray} = get_latdlonds(Grid, grid.nlat_half) + +function get_latdlonds(Grid::Type{<:AbstractGridArray}, nlat_half::Integer) + colats, lons = get_colatlons(Grid, nlat_half) # colatitudes, longitudes in radians + latds = colats # flat copy rename before conversion + londs = lons + latds .= π/2 .- colats # colatiudes to latitudes in radians + latds .*= (360/2π) # now in degrees 90˚...-90˚ + londs .*= (360/2π) + return latds, londs +end + +"""$(TYPEDSIGNATURES) Latitudes (in radians, 0-2π) and longitudes (-π/2 - π/2) for +every (horizontal) grid point in `grid`.""" +get_latlons(grid::Grid) where {Grid<:AbstractGridArray} = get_latlons(Grid, grid.nlat_half) + +function get_latlons(Grid::Type{<:AbstractGridArray}, nlat_half::Integer) + colats, lons = get_colatlons(Grid, nlat_half) # colatitudes, longitudes in radians + lats = colats # flat copy rename before conversion + lats .= π/2 .- colats # colatiudes to latitudes in radians + return lats, lons +end + +"""$(TYPEDSIGNATURES) Latitude (radians) for each ring in `grid`, north to south.""" +get_lat(grid::Grid) where {Grid<:AbstractGridArray} = get_lat(Grid, grid.nlat_half) + +"""$(TYPEDSIGNATURES) Latitude (degrees) for each ring in `grid`, north to south.""" +get_latd(grid::Grid) where {Grid<:AbstractGridArray} = get_latd(Grid, grid.nlat_half) + +"""$(TYPEDSIGNATURES) Longitude (degrees) for meridional column (full grids only).""" +get_lond(grid::Grid) where {Grid<:AbstractGridArray} = get_lond(Grid, grid.nlat_half) + +"""$(TYPEDSIGNATURES) Longitude (radians) for meridional column (full grids only).""" +get_lon(grid::Grid) where {Grid<:AbstractGridArray} = get_lon(Grid, grid.nlat_half) + +"""$(TYPEDSIGNATURES) Colatitudes (radians) for meridional column (full grids only).""" +get_colat(grid::Grid) where {Grid<:AbstractGridArray} = get_colat(Grid, grid.nlat_half) + +"""$(TYPEDSIGNATURES) Latitudes (in radians, 0-π) and longitudes (0 - 2π) for +every (horizontal) grid point in `grid`.""" +get_colatlons(grid::Grid) where {Grid<:AbstractGridArray} = get_colatlons(Grid, grid.nlat_half) +get_nlon_max(grid::Grid) where {Grid<:AbstractGridArray} = get_nlon_max(Grid, grid.nlat_half) + +function get_lat(Grid::Type{<:AbstractGridArray}, nlat_half::Integer) + return π/2 .- get_colat(Grid, nlat_half) +end + +function get_latd(Grid::Type{<:AbstractGridArray}, nlat_half::Integer) + return get_lat(Grid, nlat_half) * (360/2π) +end + +""" +$(TYPEDSIGNATURES) +Returns a vector `nlons` for the number of longitude points per latitude ring, north to south. +Provide grid `Grid` and its resolution parameter `nlat_half`. For keyword argument +`both_hemispheres=false` only the northern hemisphere (incl Equator) is returned.""" +function get_nlons(Grid::Type{<:AbstractGridArray}, nlat_half::Integer; both_hemispheres::Bool=true) + n = both_hemispheres ? get_nlat(Grid, nlat_half) : nlat_half + return [get_nlon_per_ring(Grid, nlat_half, j) for j in 1:n] +end + +## ITERATORS +""" +$(TYPEDSIGNATURES) +CartesianIndices for the 2nd to last dimension of an AbstractGridArray, +to be used like + +for k in eachgrid(grid) + for ring in eachring(grid) + for ij in ring + grid[ij, k]""" +@inline eachgrid(grid::AbstractGridArray) = CartesianIndices(size(grid)[2:end]) + +""" +$(TYPEDSIGNATURES) +Vector{UnitRange} `rings` to loop over every ring of grid `grid` +and then each grid point per ring. To be used like + + rings = eachring(grid) + for ring in rings + for ij in ring + grid[ij] + +Accesses precomputed `grid.rings`.""" +@inline eachring(grid::AbstractGridArray) = grid.rings + +"""$(TYPEDSIGNATURES) +Computes the ring indices `i0:i1` for start and end of every longitudinal point +on a given ring `j` of `Grid` at resolution `nlat_half`. Used to loop +over rings of a grid. These indices are also precomputed in every `grid.rings`.""" +function eachring(Grid::Type{<:AbstractGridArray}, nlat_half::Integer) + rings = Vector{UnitRange{Int}}(undef, get_nlat(Grid, nlat_half)) # allocate + each_index_in_ring!(rings, Grid, nlat_half) # calculate iteratively + return rings +end + +"""$(TYPEDSIGNATURES) Same as `eachring(grid)` but performs a bounds check to assess +that all `grids` are of same size.""" +function eachring(grid1::Grid, grids::Grid...) where {Grid<:AbstractGridArray} + n = length(grid1) + Base._all_match_first(X->length(X), n, grid1, grids...) || throw(BoundsError) + return eachring(grid1) +end + +# equality and comparison, somehow needed as not covered by broadcasting +Base.:(==)(G1::AbstractGridArray, G2::AbstractGridArray) = grids_match(G1, G2) && G1.data == G2.data +Base.all(G::AbstractGridArray) = all(G.data) +Base.any(G::AbstractGridArray) = any(G.data) + +"""$(TYPEDSIGNATURES) True if both `A` and `B` are of the same type +(regardless type parameter `T` or underyling array type `ArrayType`) and +of same size.""" +function grids_match(A::AbstractGridArray, B::AbstractGridArray) + length(A) == length(B) && return grids_match(typeof(A), typeof(B)) + return false +end + +function grids_match(A::Type{<:AbstractGridArray}, B::Type{<:AbstractGridArray}) + # eltypes can be different and also array types of underlying data + return nonparametric_type(A) == nonparametric_type(B) +end + +"""$(TYPEDSIGNATURES) UnitRange to access data on grid `grid` on ring `j`.""" +function each_index_in_ring(grid::Grid, j::Integer) where {Grid<:AbstractGridArray} + return each_index_in_ring(Grid, j, grid.nlat_half) +end + +""" $(TYPEDSIGNATURES) UnitRange to access each grid point on grid `grid`.""" +eachgridpoint(grid::AbstractGridArray) = Base.OneTo(get_npoints(grid)) +function eachgridpoint(grid1::Grid, grids::Grid...) where {Grid<:AbstractGridArray} + n = length(grid1) + Base._all_match_first(X->length(X), n, grid1, grids...) || throw(BoundsError) + return eachgridpoint(grid1) +end + +"""$(TYPEDSIGNATURES) Obtain ring index `j` from gridpoint `ij` and `rings` +describing rind indices as obtained from `eachring(::Grid)`""" +function whichring(ij::Integer, rings::Vector{UnitRange{Int}}) + @boundscheck 0 < ij <= rings[end][end] || throw(BoundsError) + j = 1 + @inbounds while ij > rings[j][end] + j += 1 + end + return j +end + +## BROADCASTING +# following https://docs.julialang.org/en/v1/manual/interfaces/#man-interfaces-broadcasting +import Base.Broadcast: BroadcastStyle, Broadcasted, DefaultArrayStyle +import LinearAlgebra: isstructurepreserving, fzeropreserving + +# {1} as grids are <:AbstractVector, Grid here is the non-parameteric Grid type! +struct AbstractGridArrayStyle{N, Grid} <: Broadcast.AbstractArrayStyle{N} end + +# important to remove Grid{T} parameter T (eltype/number format) here to broadcast +# automatically across the same grid type but with different T +# e.g. FullGaussianGrid{Float32} and FullGaussianGrid{Float64} +Base.BroadcastStyle(::Type{Grid}) where {Grid<:AbstractGridArray{T, N, ArrayType}} where {T, N, ArrayType} = + AbstractGridArrayStyle{N, nonparametric_type(Grid)}() + +# allocation for broadcasting, create a new Grid with undef of type/number format T +function Base.similar(bc::Broadcasted{AbstractGridArrayStyle{N, Grid}}, ::Type{T}) where {N, Grid, T} + return Grid(Array{T}(undef, size(bc)...)) +end + +# ::Val{0} for broadcasting with 0-dimensional, ::Val{1} for broadcasting with vectors, etc +AbstractGridArrayStyle{N, Grid}(::Val{M}) where {N, Grid, M} = AbstractGridArrayStyle{N, Grid}() + +## GPU +struct AbstractGPUGridArrayStyle{N, ArrayType, Grid} <: GPUArrays.AbstractGPUArrayStyle{N} end + +function Base.BroadcastStyle( + ::Type{Grid} +) where {Grid<:AbstractGridArray{T, N, ArrayType}} where {T, N, ArrayType <: GPUArrays.AbstractGPUArray} + return AbstractGPUGridArrayStyle{N, ArrayType, nonparametric_type(Grid)}() +end + +# ::Val{0} for broadcasting with 0-dimensional, ::Val{1} for broadcasting with vectors, etc +AbstractGPUGridArrayStyle{N, ArrayType, Grid}(::Val{M}) where {N, ArrayType, Grid, M} = + AbstractGPUGridArrayStyle{N, ArrayType, Grid}() + +function GPUArrays.backend( + ::Type{Grid} +) where {Grid <: AbstractGridArray{T, N, ArrayType}} where {T, N, ArrayType <: GPUArrays.AbstractGPUArray} + return GPUArrays.backend(ArrayType) +end + +function Base.similar( + bc::Broadcasted{AbstractGPUGridArrayStyle{N, ArrayType, Grid}}, + ::Type{T}, +) where {N, ArrayType, Grid, T} + ArrayType_ = nonparametric_type(ArrayType) + return Grid(ArrayType_{T}(undef, size(bc)...)) +end + +function Adapt.adapt_structure(to, grid::Grid) where {Grid <: AbstractGridArray} + Grid_ = nonparametric_type(Grid) + return Grid_(Adapt.adapt(to, grid.data), grid.nlat_half, grid.rings) +end \ No newline at end of file diff --git a/src/RingGrids/grids/full_clenshaw.jl b/src/RingGrids/grids/full_clenshaw.jl new file mode 100644 index 000000000..2e10643f4 --- /dev/null +++ b/src/RingGrids/grids/full_clenshaw.jl @@ -0,0 +1,41 @@ +"""A `FullClenshawArray` is an array of full grid, subtyping `AbstractFullGridArray`, that use +equidistant latitudes for each ring (a regular lon-lat grid). These require the +Clenshaw-Curtis quadrature in the spectral transform, hence the name. One ring is on the equator, +total number of rings is odd, no rings on the north or south pole. First dimension of the +underlying `N`-dimensional `data` represents the horizontal dimension, in ring order +(0 to 360˚E, then north to south), other dimensions are used for the vertical and/or +time or other dimensions. The resolution parameter of the horizontal grid is `nlat_half` +(number of latitude rings on one hemisphere, Equator included) and the ring indices are +precomputed in `rings`. Fields are +$(TYPEDFIELDS)""" +struct FullClenshawArray{T, N, ArrayType <: AbstractArray{T, N}} <: AbstractFullGridArray{T, N, ArrayType} + data::ArrayType # data array, ring by ring, north to south + nlat_half::Int # number of latitudes on one hemisphere + rings::Vector{UnitRange{Int}} # TODO make same array type as data? + + FullClenshawArray(data::A, nlat_half, rings) where {A <: AbstractArray{T, N}} where {T, N} = + check_inputs(data, nlat_half, rings, FullClenshawArray) ? + new{T, N, A}(data, nlat_half, rings) : + error_message(data, nlat_half, rings, FullClenshawArray, T, N, A) +end + +# TYPES +const FullClenshawGrid{T} = FullClenshawArray{T, 1, Vector{T}} +nonparametric_type(::Type{<:FullClenshawArray}) = FullClenshawArray +horizontal_grid_type(::Type{<:FullClenshawArray}) = FullClenshawGrid + +"""A `FullClenshawArray` but constrained to `N=1` dimensions (horizontal only) and data is a `Vector{T}`.""" +FullClenshawGrid + +# SIZE +nlat_odd(::Type{<:FullClenshawArray}) = true +get_npoints2D(::Type{<:FullClenshawArray}, nlat_half::Integer) = 8 * nlat_half^2 - 4nlat_half +get_nlat_half(::Type{<:FullClenshawArray}, npoints2D::Integer) = round(Int, 1/4 + sqrt(1/16 + npoints2D/8)) +get_nlon(::Type{<:FullClenshawArray}, nlat_half::Integer) = 4nlat_half + +## COORDINATES +get_colat(::Type{<:FullClenshawArray}, nlat_half::Integer) = [j/(2nlat_half)*π for j in 1:2nlat_half-1] +get_lon(::Type{<:FullClenshawArray}, nlat_half::Integer) = get_lon(FullGaussianArray, nlat_half) + +# QUADRATURE +get_quadrature_weights(::Type{<:FullClenshawArray}, nlat_half::Integer) = clenshaw_curtis_weights(nlat_half) \ No newline at end of file diff --git a/src/RingGrids/grids/full_gaussian.jl b/src/RingGrids/grids/full_gaussian.jl new file mode 100644 index 000000000..044aed48c --- /dev/null +++ b/src/RingGrids/grids/full_gaussian.jl @@ -0,0 +1,51 @@ +"""A `FullGaussianArray` is an array of full grids, subtyping `AbstractFullGridArray`, that use +Gaussian latitudes for each ring. First dimension of the underlying `N`-dimensional `data` +represents the horizontal dimension, in ring order (0 to 360˚E, then north to south), +other dimensions are used for the vertical and/or time or other dimensions. +The resolution parameter of the horizontal grid is `nlat_half` (number of latitude rings +on one hemisphere, Equator included) and the ring indices are precomputed in `rings`. +Fields are +$(TYPEDFIELDS)""" +struct FullGaussianArray{T, N, ArrayType <: AbstractArray{T, N}} <: AbstractFullGridArray{T, N, ArrayType} + "Data array, west to east, ring by ring, north to south." + data::ArrayType + + "Number of latitudes on one hemisphere" + nlat_half::Int + + "Precomputed ring indices, ranging from first to last grid point on every ring." + rings::Vector{UnitRange{Int}} # TODO make same array type as data? + + FullGaussianArray(data::A, nlat_half, rings) where {A <: AbstractArray{T, N}} where {T, N} = + check_inputs(data, nlat_half, rings, FullGaussianArray) ? + new{T, N, A}(data, nlat_half, rings) : + error_message(data, nlat_half, rings, FullGaussianArray, T, N, A) +end + +# TYPES +const FullGaussianGrid{T} = FullGaussianArray{T, 1, Vector{T}} +nonparametric_type(::Type{<:FullGaussianArray}) = FullGaussianArray # identity for full grids +horizontal_grid_type(::Type{<:FullGaussianArray}) = FullGaussianGrid + +"""A `FullGaussianArray` but constrained to `N=1` dimensions (horizontal only) and data is a `Vector{T}`.""" +FullGaussianGrid + +# SIZE +nlat_odd(::Type{<:FullGaussianArray}) = false # Gaussian latitudes always even +get_npoints2D(::Type{<:FullGaussianArray}, nlat_half::Integer) = 8 * nlat_half^2 +get_nlat_half(::Type{<:FullGaussianArray}, npoints2D::Integer) = round(Int, sqrt(npoints2D/8)) +get_nlon(::Type{<:FullGaussianArray}, nlat_half::Integer) = 4nlat_half + +## COORDINATES +function get_colat(::Type{<:FullGaussianArray}, nlat_half::Integer) + return π .- acos.(FastGaussQuadrature.gausslegendre(2nlat_half)[1]) +end + +function get_lon(::Type{<:FullGaussianArray}, nlat_half::Integer) + nlat_half == 0 && return Float64[] + nlon = get_nlon(FullGaussianArray, nlat_half) + return collect(range(0, 2π-π/nlon, step=2π/nlon)) +end + +# QUADRATURE +get_quadrature_weights(::Type{<:FullGaussianArray}, nlat_half::Integer) = gaussian_weights(nlat_half) \ No newline at end of file diff --git a/src/RingGrids/grids/full_healpix.jl b/src/RingGrids/grids/full_healpix.jl new file mode 100644 index 000000000..8a6aefed4 --- /dev/null +++ b/src/RingGrids/grids/full_healpix.jl @@ -0,0 +1,41 @@ +"""A `FullHEALPixArray` is an array of full grids, subtyping `AbstractFullGridArray`, that use +HEALPix latitudes for each ring. This type primarily equists to interpolate data from +the (reduced) HEALPixGrid onto a full grid for output. + +First dimension of the underlying `N`-dimensional `data` represents the horizontal dimension, +in ring order (0 to 360˚E, then north to south), other dimensions are used for the vertical +and/or time or other dimensions. The resolution parameter of the horizontal grid is +`nlat_half` (number of latitude rings on one hemisphere, Equator included) and the ring indices +are precomputed in `rings`. Fields are +$(TYPEDFIELDS)""" +struct FullHEALPixArray{T, N, ArrayType <: AbstractArray{T, N}} <: AbstractFullGridArray{T, N, ArrayType} + data::ArrayType # data array, ring by ring, north to south + nlat_half::Int # number of latitudes on one hemisphere + rings::Vector{UnitRange{Int}} # TODO make same array type as data? + + FullHEALPixArray(data::A, nlat_half, rings) where {A <: AbstractArray{T, N}} where {T, N} = + check_inputs(data, nlat_half, rings, FullHEALPixArray) ? + new{T, N, A}(data, nlat_half, rings) : + error_message(data, nlat_half, rings, FullHEALPixArray, T, N, A) +end + +# TYPES +const FullHEALPixGrid{T} = FullHEALPixArray{T, 1, Vector{T}} +nonparametric_type(::Type{<:FullHEALPixArray}) = FullHEALPixArray +horizontal_grid_type(::Type{<:FullHEALPixArray}) = FullHEALPixGrid + +"""A `FullHEALPixArray` but constrained to `N=1` dimensions (horizontal only) and data is a `Vector{T}`.""" +FullHEALPixGrid + +# SIZE +nlat_odd(::Type{<:FullHEALPixArray}) = true +get_npoints2D(::Type{<:FullHEALPixArray}, nlat_half::Integer) = 4nlat_half * (2nlat_half-1) +get_nlat_half(::Type{<:FullHEALPixArray}, npoints2D::Integer) = round(Int, 1/4 + sqrt(1/16 + npoints2D/8)) +get_nlon(::Type{<:FullHEALPixArray}, nlat_half::Integer) = 4nlat_half + +## COORDINATES +get_colat(::Type{<:FullHEALPixArray}, nlat_half::Integer) = get_colat(HEALPixGrid, nlat_half) +get_lon(::Type{<:FullHEALPixArray}, nlat_half::Integer) = get_lon(FullGaussianArray, nlat_half) + +# QUADRATURE (use weights from reduced grids though!) +get_quadrature_weights(::Type{<:FullHEALPixArray}, nlat_half::Integer) = equal_area_weights(HEALPixArray, nlat_half) \ No newline at end of file diff --git a/src/RingGrids/grids/full_octahealpix.jl b/src/RingGrids/grids/full_octahealpix.jl new file mode 100644 index 000000000..1e604bfcf --- /dev/null +++ b/src/RingGrids/grids/full_octahealpix.jl @@ -0,0 +1,42 @@ +"""A `FullOctaHEALPixArray` is an array of full grids, subtyping `AbstractFullGridArray` that use +OctaHEALPix latitudes for each ring. This type primarily equists to interpolate data from +the (reduced) OctaHEALPixGrid onto a full grid for output. + +First dimension of the underlying `N`-dimensional `data` represents the horizontal dimension, +in ring order (0 to 360˚E, then north to south), other dimensions are used for the vertical +and/or time or other dimensions. The resolution parameter of the horizontal grid is +`nlat_half` (number of latitude rings on one hemisphere, Equator included) and the ring indices +are precomputed in `rings`. Fields are +$(TYPEDFIELDS)""" +struct FullOctaHEALPixArray{T, N, ArrayType <: AbstractArray{T, N}} <: AbstractFullGridArray{T, N, ArrayType} + data::ArrayType # data array, ring by ring, north to south + nlat_half::Int # number of latitudes on one hemisphere + rings::Vector{UnitRange{Int}} # TODO make same array type as data? + + FullOctaHEALPixArray(data::A, nlat_half, rings) where {A <: AbstractArray{T, N}} where {T, N} = + check_inputs(data, nlat_half, rings, FullOctaHEALPixArray) ? + new{T, N, A}(data, nlat_half, rings) : + error_message(data, nlat_half, rings, FullOctaHEALPixArray, T, N, A) +end + +# TYPES +const FullOctaHEALPixGrid{T} = FullOctaHEALPixArray{T, 1, Vector{T}} +nonparametric_type(::Type{<:FullOctaHEALPixArray}) = FullOctaHEALPixArray +horizontal_grid_type(::Type{<:FullOctaHEALPixArray}) = FullOctaHEALPixGrid + +"""A `FullOctaHEALPixArray` but constrained to `N=1` dimensions (horizontal only) and data is a `Vector{T}`.""" +FullOctaHEALPixGrid + +# SIZE +nlat_odd(::Type{<:FullOctaHEALPixArray}) = true +get_npoints2D(::Type{<:FullOctaHEALPixArray}, nlat_half::Integer) = 4nlat_half * (2nlat_half-1) +get_nlat_half(::Type{<:FullOctaHEALPixArray}, npoints2D::Integer) = round(Int, 1/4 + sqrt(1/16 + npoints2D/8)) +get_nlon(::Type{<:FullOctaHEALPixArray}, nlat_half::Integer) = 4nlat_half + +## COORDINATES +get_colat(::Type{<:FullOctaHEALPixArray}, nlat_half::Integer) = get_colat(OctaHEALPixGrid, nlat_half) +get_lon(::Type{<:FullOctaHEALPixArray}, nlat_half::Integer) = get_lon(FullGaussianArray, nlat_half) + +# QUADRATURE (use weights from reduced grids though!) +get_quadrature_weights(::Type{<:FullOctaHEALPixArray}, nlat_half::Integer) = + equal_area_weights(OctaHEALPixArray, nlat_half) \ No newline at end of file diff --git a/src/RingGrids/grids/healpix.jl b/src/RingGrids/grids/healpix.jl new file mode 100644 index 000000000..c8f871495 --- /dev/null +++ b/src/RingGrids/grids/healpix.jl @@ -0,0 +1,158 @@ +"""A `HEALPixArray` is an array of HEALPix grids, subtyping `AbstractReducedGridArray`. +First dimension of the underlying `N`-dimensional `data` represents the horizontal dimension, +in ring order (0 to 360˚E, then north to south), other dimensions are used for the vertical and/or +time or other dimensions. The resolution parameter of the horizontal grid is `nlat_half` +(number of latitude rings on one hemisphere, Equator included) which has to be even +(non-fatal error thrown otherwise) which is less strict than the original HEALPix +formulation (only powers of two for nside = nlat_half/2). Ring indices are precomputed in `rings`. + +A HEALPix grid has 12 faces, each `nside`x`nside` grid points, each covering the same area +of the sphere. They start with 4 longitude points on the northern-most ring, +increase by 4 points per ring in the "polar cap" (the top half of the 4 northern-most faces) +but have a constant number of longitude points in the equatorial belt. The southern hemisphere +is symmetric to the northern, mirrored around the Equator. HEALPix grids have a ring on the +Equator. For more details see Górski et al. 2005, DOI:10.1086/427976. + +`rings` are the precomputed ring indices, for nlat_half = 4 it is +`rings = [1:4, 5:12, 13:20, 21:28, 29:36, 37:44, 45:48]`. So the first ring has indices +1:4 in the unravelled first dimension, etc. For efficient looping see `eachring` and `eachgrid`. +Fields are +$(TYPEDFIELDS)""" +struct HEALPixArray{T, N, ArrayType <: AbstractArray{T, N}} <: AbstractReducedGridArray{T, N, ArrayType} + data::ArrayType # data array, ring by ring, north to south + nlat_half::Int # number of latitudes on one hemisphere + rings::Vector{UnitRange{Int}} # TODO make same array type as data? + + HEALPixArray(data::A, nlat_half, rings) where {A <: AbstractArray{T, N}} where {T, N} = + check_inputs(data, nlat_half, rings, HEALPixArray) ? + new{T, N, A}(data, nlat_half, rings) : + error_message(data, nlat_half, rings, HEALPixArray, T, N, A) +end + +## TYPES +const HEALPixGrid{T} = HEALPixArray{T, 1, Vector{T}} +nonparametric_type(::Type{<:HEALPixArray}) = HEALPixArray +horizontal_grid_type(::Type{<:HEALPixArray}) = HEALPixGrid +full_array_type(::Type{<:HEALPixArray}) = FullHEALPixArray + +"""An `HEALPixArray` but constrained to `N=1` dimensions (horizontal only) and data is a `Vector{T}`.""" +HEALPixGrid + +## SIZE +nlat_odd(::Type{<:HEALPixArray}) = true +get_npoints2D(::Type{<:HEALPixArray}, nlat_half::Integer) = 3*nlat_half^2 +get_nlat_half(::Type{<:HEALPixArray}, npoints2D::Integer) = round(Int, sqrt(npoints2D/3)) + +"""$(TYPEDSIGNATURES) Number of longitude points for ring `j` on `Grid` of resolution +`nlat_half`.""" +function get_nlon_per_ring(Grid::Type{<:HEALPixArray}, nlat_half::Integer, j::Integer) + nlat = get_nlat(Grid, nlat_half) + @assert 0 < j <= nlat "Ring $j is outside H$nlat_half grid." + return min(4j, 2nlat_half, 8nlat_half-4j) +end + +"""$(TYPEDSIGNATURES) The original `Nside` resolution parameter of the HEALPix grids. +The number of grid points on one side of each (square) face. +While we use `nlat_half` across all ring grids, this function translates this to +Nside. Even `nlat_half` only.""" +function nside_healpix(nlat_half::Integer) + iseven(nlat_half) || @error "Odd nlat_half=$nlat_half not supported for HEALPix." + return nlat_half÷2 +end + +# for future reordering the HEALPix ring order into a matrix consisting of the +# 12 square faces of a HEALPix grid. +function matrix_size(::Type{<:HEALPixArray}, nlat_half::Integer) + nside = nside_healpix(nlat_half) + return (5nside, 5nside) +end + +## COORDINATES +function get_colat(::Type{<:HEALPixArray}, nlat_half::Integer) + nlat_half == 0 && return Float64[] + + nlat = get_nlat(HEALPixArray, nlat_half) + nside = nside_healpix(nlat_half) + colat = zeros(nlat) + + # Górski et al 2005 eq 4 and 8 + for j in 1:nside colat[j] = acos(1-j^2/3nside^2) end # north polar cap + for j in nside+1:3nside colat[j] = acos(4/3 - 2j/3nside) end # equatorial belt + for j in 3nside+1:nlat colat[j] = acos((2nlat_half-j)^2/3nside^2-1) end # south polar cap + + return colat +end + +function get_lon_per_ring(Grid::Type{<:HEALPixArray}, nlat_half::Integer, j::Integer) + nside = nside_healpix(nlat_half) + nlon = get_nlon_per_ring(Grid, nlat_half, j) + + # Górski et al 2005 eq 5 and 9 + # s = 1 for polar caps, s=2, 1, 2, 1, ... in the equatorial zone + s = (j < nside) || (j >= 3nside) ? 1 : ((j - nside) % 2 + 1) + lon = [π/(nlon÷2)*(i - s/2) for i in 1:nlon] + return lon +end + +## INDEXING +function each_index_in_ring(::Type{<:HEALPixArray}, + j::Integer, # ring index north to south + nlat_half::Integer) # resolution param + nside = nside_healpix(nlat_half) + + @boundscheck 0 < j < 4nside || throw(BoundsError) # ring index valid? + if j < nside # northern polar cap + index_1st = 2j*(j-1) + 1 # first in-ring index i + index_end = 2j*(j+1) # last in-ring index i + elseif j <= 3nside # equatorial zone with const nlon + n = 2nside^2-2nside # points in polar cap + nlon = 4nside # points on latitude rings + j = j - nside + 1 # offset ring index into eq zone + index_1st = n + (j-1)*nlon + 1 # add constant nlon per ring + index_end = n + j*nlon + else # southern polar cap + n = 12nside^2 # total number of points + j = 4nside - j # count ring index from south pole + index_1st = n - 2j*(j+1) + 1 # count backwards + index_end = n - 2j*(j-1) + end + return index_1st:index_end # range of i's in ring +end + +function each_index_in_ring!( rings::Vector{<:UnitRange{<:Integer}}, + Grid::Type{<:HEALPixArray}, + nlat_half::Integer) # resolution param + + nlat = length(rings) + @boundscheck nlat == get_nlat(Grid, nlat_half) || throw(BoundsError) + + # HEALPix not defined for odd nlat_half, the last rings would not be written + @boundscheck iseven(nlat_half) || throw(BoundsError) + + index_end = 0 + nside = nside_healpix(nlat_half) # side length of a basepixel + + # North polar cap + @inbounds for j in 1:nside-1 + index_1st = index_end + 1 # 1st index is +1 from prev ring's last index + index_end += 4j # add number of grid points per ring + rings[j] = index_1st:index_end # turn into UnitRange + end + + # Equatorial belt + nlon_max = get_nlon_max(Grid, nlat_half) # number of grid points on belt + @inbounds for j in nside:3nside + index_1st = index_end + 1 # 1st index is +1 from prev ring's last index + index_end += nlon_max # nlon constant in belt + rings[j] = index_1st:index_end # turn into UnitRange + end + + # South polar cap + @inbounds for (j, j_mirrored) in zip( 3nside+1:nlat, # South only + nside-1:-1:1) # mirror index + + index_1st = index_end + 1 # 1st index is +1 from prev ring's last index + index_end += 4j_mirrored # add number of grid points per ring + rings[j] = index_1st:index_end # turn into UnitRange + end +end \ No newline at end of file diff --git a/src/RingGrids/octahealpix.jl b/src/RingGrids/grids/octahealpix.jl similarity index 62% rename from src/RingGrids/octahealpix.jl rename to src/RingGrids/grids/octahealpix.jl index a2c5f344f..bcdf5b1eb 100644 --- a/src/RingGrids/octahealpix.jl +++ b/src/RingGrids/grids/octahealpix.jl @@ -1,47 +1,80 @@ -""" - abstract type AbstractOctaHEALPixGrid{T} <: AbstractGrid{T} end - -An `AbstractOctaHEALPixGrid` is a horizontal grid similar to the standard OctahedralGrid, -but the number of points in the ring closest to the Poles starts from 4 instead of 20, -and the longitude of the first point in each ring is shifted as in HEALPixGrid. -Also, different latitudes can be used.""" -abstract type AbstractOctaHEALPixGrid{T} <: AbstractGrid{T} end - -nlat_odd(::Type{<:AbstractOctaHEALPixGrid}) = true -get_nlon_max(::Type{<:AbstractOctaHEALPixGrid}, nlat_half::Integer) = 4nlat_half - -function get_nlon_per_ring(Grid::Type{<:AbstractOctaHEALPixGrid}, nlat_half::Integer, j::Integer) - nlat = get_nlat(Grid, nlat_half) - @assert 0 < j <= nlat "Ring $j is outside H$nlat_half grid." - j = j > nlat_half ? nlat - j + 1 : j # flip north south due to symmetry - return nlon_octahealpix(nlat_half, j) +"""An `OctaHEALPixArray` is an array of OctaHEALPix grids, subtyping `AbstractReducedGridArray`. +First dimension of the underlying `N`-dimensional `data` represents the horizontal dimension, +in ring order (0 to 360˚E, then north to south), other dimensions are used for the vertical and/or +time or other dimensions. The resolution parameter of the horizontal grid is `nlat_half` +(number of latitude rings on one hemisphere, Equator included) and the ring indices are +precomputed in `rings`. + +An OctaHEALPix grid has 4 faces, each `nlat_half x nlat_half` in size, +covering 90˚ in longitude, pole to pole. As part of the HEALPix family of grids, +the grid points are equal area. They start with 4 longitude points on the northern-most ring, +increase by 4 points per ring towards the Equator with one ring on the Equator before reducing +the number of points again towards the south pole by 4 per ring. There is no equatorial belt for +OctaHEALPix grids. The southern hemisphere is symmetric to the northern, mirrored around the Equator. +OctaHEALPix grids have a ring on the Equator. For more details see +Górski et al. 2005, DOI:10.1086/427976, the OctaHEALPix grid belongs to the family of +HEALPix grids with Nθ = 1, Nφ = 4 but is not explicitly mentioned therein. + +`rings` are the precomputed ring indices, for nlat_half = 3 (in contrast to HEALPix this can be odd) +it is `rings = [1:4, 5:12, 13:24, 25:32, 33:36]`. For efficient looping see `eachring` and `eachgrid`. +Fields are +$(TYPEDFIELDS)""" +struct OctaHEALPixArray{T, N, ArrayType <: AbstractArray{T, N}} <: AbstractReducedGridArray{T, N, ArrayType} + data::ArrayType # data array, ring by ring, north to south + nlat_half::Int # number of latitudes on one hemisphere + rings::Vector{UnitRange{Int}} # TODO make same array type as data? + + OctaHEALPixArray(data::A, nlat_half, rings) where {A <: AbstractArray{T, N}} where {T, N} = + check_inputs(data, nlat_half, rings, OctaHEALPixArray) ? + new{T, N, A}(data, nlat_half, rings) : + error_message(data, nlat_half, rings, OctaHEALPixArray, T, N, A) end -get_npoints(::Type{<:AbstractOctaHEALPixGrid}, nlat_half::Integer) = npoints_octahealpix(nlat_half) -get_lon(::Type{<:AbstractOctaHEALPixGrid}, nlat_half::Integer) = Float64[] # only defined for full grids +## TYPES +const OctaHEALPixGrid{T} = OctaHEALPixArray{T, 1, Vector{T}} +nonparametric_type(::Type{<:OctaHEALPixArray}) = OctaHEALPixArray +horizontal_grid_type(::Type{<:OctaHEALPixArray}) = OctaHEALPixGrid +full_array_type(::Type{<:OctaHEALPixArray}) = FullOctaHEALPixArray -function get_colatlons(Grid::Type{<:AbstractOctaHEALPixGrid}, nlat_half::Integer) - nlat = get_nlat(Grid, nlat_half) - npoints = get_npoints(Grid, nlat_half) - colat = get_colat(Grid, nlat_half) +"""An `OctaHEALPixArray` but constrained to `N=1` dimensions (horizontal only) and data is a `Vector{T}`.""" +OctaHEALPixGrid - colats = zeros(npoints) - lons = zeros(npoints) +## SIZE +nlat_odd(::Type{<:OctaHEALPixArray}) = true +get_npoints2D(::Type{<:OctaHEALPixArray}, nlat_half::Integer) = 4*nlat_half^2 +get_nlat_half(::Type{<:OctaHEALPixArray}, npoints2D::Integer) = round(Int, sqrt(npoints2D/4)) - ij = 1 - for j in 1:nlat - nlon = get_nlon_per_ring(Grid, nlat_half, j) - lon = collect(π/nlon:2π/nlon:2π) +# number of longitude +function get_nlon_per_ring(Grid::Type{<:OctaHEALPixArray}, nlat_half::Integer, j::Integer) + nlat = get_nlat(Grid, nlat_half) + @assert 0 < j <= nlat "Ring $j is outside P$nlat_half grid." + # j = j > nlat_half ? nlat - j + 1 : j # flip north south due to symmetry + return min(4j, 8nlat_half-4j) +end - colats[ij:ij+nlon-1] .= colat[j] - lons[ij:ij+nlon-1] .= lon +matrix_size(::Type{OctaHEALPixGrid}, nlat_half::Integer) = (2nlat_half, 2nlat_half) - ij += nlon +## COORDINATES +function get_colat(::Type{<:OctaHEALPixArray}, nlat_half::Integer) + nlat_half == 0 && return Float64[] + colat = zeros(get_nlat(OctaHEALPixArray, nlat_half)) + for j in 1:nlat_half + # Górski et al. 2005 eq 4 but without the 1/3 and Nside=nlat_half + colat[j] = acos(1-(j/nlat_half)^2) # northern hemisphere + colat[2nlat_half-j] = π - colat[j] # southern hemisphere end - return colats, lons + return colat +end + +function get_lon_per_ring(Grid::Type{<:OctaHEALPixArray}, nlat_half::Integer, j::Integer) + nlon = get_nlon_per_ring(Grid, nlat_half, j) + # equidistant longitudes with equal offsets from 0˚ and 360˚, + # e.g. 45, 135, 225, 315 for nlon=4 + return collect(π/nlon:2π/nlon:2π) end -function each_index_in_ring(::Type{<:AbstractOctaHEALPixGrid}, # function for OctaHEALPix grids +## INDEXING +function each_index_in_ring(::Type{<:OctaHEALPixArray}, # function for OctaHEALPix grids j::Integer, # ring index north to south nlat_half::Integer) # resolution param @@ -59,7 +92,7 @@ function each_index_in_ring(::Type{<:AbstractOctaHEALPixGrid}, # function for Oc end function each_index_in_ring!( rings::Vector{<:UnitRange{<:Integer}}, - Grid::Type{<:AbstractOctaHEALPixGrid}, + Grid::Type{<:OctaHEALPixArray}, nlat_half::Integer) # resolution param nlat = length(rings) @@ -71,7 +104,7 @@ function each_index_in_ring!( rings::Vector{<:UnitRange{<:Integer}}, index_end += 4j # add number of grid points per ring rings[j] = index_1st:index_end # turn into UnitRange end - @inbounds for (j, j_rev) in zip( nlat_half+1:nlat, # South only + @inbounds for (j, j_rev) in zip(nlat_half+1:nlat, # South only nlat-nlat_half:-1:1) # reverse index index_1st = index_end + 1 # 1st index is +1 from prev ring's last index @@ -80,47 +113,7 @@ function each_index_in_ring!( rings::Vector{<:UnitRange{<:Integer}}, end end - -""" - H = OctaHEALPixGrid{T} - -A OctaHEALPix grid with 4 base faces, each `nlat_half`x`nlat_half` grid points, each covering the same area. -The values of all grid points are stored in a vector field `data` that unravels the data 0 to 360˚, -then ring by ring, which are sorted north to south.""" -struct OctaHEALPixGrid{T} <: AbstractOctaHEALPixGrid{T} - data::Vector{T} # data vector, ring by ring, north to south - nlat_half::Int # number of latitude rings on one hemisphere - - OctaHEALPixGrid{T}(data::AbstractVector, nlat_half::Integer) where T = length(data) == npoints_octahealpix(nlat_half) ? - new(data, nlat_half) : error("$(length(data))-element Vector{$(eltype(data))}"* - "cannot be used to create an H$nlat_half OctaHEALPixGrid{$T}.") -end - -nonparametric_type(::Type{<:OctaHEALPixGrid}) = OctaHEALPixGrid - -npoints_octahealpix(nlat_half::Integer) = 4nlat_half^2 -nlat_half_octahealpix(npoints::Integer) = round(Int, sqrt(npoints/4)) # inverse of npoints_octahealpix -nlat_octahealpix(nlat_half::Integer) = 2nlat_half-1 -nlon_octahealpix(nlat_half::Integer, j::Integer) = min(4j, 8nlat_half-4j) - -# infer nlat_half from data vector length, infer parametric type from eltype of data -OctaHEALPixGrid{T}(data::AbstractVector) where T = OctaHEALPixGrid{T}(data, nlat_half_octahealpix(length(data))) -OctaHEALPixGrid(data::AbstractVector, n::Integer...) = OctaHEALPixGrid{eltype(data)}(data, n...) - -function get_colat(::Type{<:OctaHEALPixGrid}, nlat_half::Integer) - nlat_half == 0 && return Float64[] - colat = zeros(nlat_octahealpix(nlat_half)) - for j in 1:nlat_half - colat[j] = acos(1-(j/nlat_half)^2) # northern hemisphere - colat[2nlat_half-j] = π - colat[j] # southern hemisphere - end - return colat -end - -full_grid(::Type{<:OctaHEALPixGrid}) = FullOctaHEALPixGrid # the full grid with same latitudes - -matrix_size(grid::OctaHEALPixGrid) = (2grid.nlat_half, 2grid.nlat_half) -matrix_size(::Type{OctaHEALPixGrid}, nlat_half::Integer) = (2nlat_half, 2nlat_half) +## CONVERSION Base.Matrix(G::OctaHEALPixGrid{T}; kwargs...) where T = Matrix!(zeros(T, matrix_size(G)...), G; kwargs...) """ diff --git a/src/RingGrids/grids/octahedral_clenshaw.jl b/src/RingGrids/grids/octahedral_clenshaw.jl new file mode 100644 index 000000000..1314d6c54 --- /dev/null +++ b/src/RingGrids/grids/octahedral_clenshaw.jl @@ -0,0 +1,210 @@ +"""An `OctahedralClenshawArray` is an array of octahedral grids, subtyping `AbstractReducedGridArray`, +that use equidistant latitudes for each ring, the same as for `FullClenshawArray`. +First dimension of the underlying `N`-dimensional `data` represents the horizontal dimension, +in ring order (0 to 360˚E, then north to south), other dimensions are used for the vertical and/or +time or other dimensions. The resolution parameter of the horizontal grid is `nlat_half` +(number of latitude rings on one hemisphere, Equator included) and the ring indices are +precomputed in `rings`. + +These grids are called octahedral (same as for the `OctahedralGaussianArray` which only uses different +latitudes) because after starting with 20 points on the first ring around the north pole (default) they +increase the number of longitude points for each ring by 4, such that they can be conceptually thought +of as lying on the 4 faces of an octahedron on each hemisphere. Hence, these grids have 20, 24, 28, ... +longitude points for ring 1, 2, 3, ... Clenshaw grids have a ring on the Equator which has 16 + 4nlat_half +longitude points before reducing the number of longitude points per ring by 4 towards the southern-most +ring `j = nlat`. `rings` are the precomputed ring indices, the the example above +`rings = [1:20, 21:44, 45:72, ...]`. For efficient looping see `eachring` and `eachgrid`. +Fields are +$(TYPEDFIELDS)""" +struct OctahedralClenshawArray{T, N, ArrayType <: AbstractArray{T, N}} <: AbstractReducedGridArray{T, N, ArrayType} + data::ArrayType # data array, ring by ring, north to south + nlat_half::Int # number of latitudes on one hemisphere + rings::Vector{UnitRange{Int}} # TODO make same array type as data? + + OctahedralClenshawArray(data::A, nlat_half, rings) where {A <: AbstractArray{T, N}} where {T, N} = + check_inputs(data, nlat_half, rings, OctahedralClenshawArray) ? + new{T, N, A}(data, nlat_half, rings) : + error_message(data, nlat_half, rings, OctahedralClenshawArray, T, N, A) +end + +# TYPES +const OctahedralClenshawGrid{T} = OctahedralClenshawArray{T, 1, Vector{T}} +nonparametric_type(::Type{<:OctahedralClenshawArray}) = OctahedralClenshawArray +horizontal_grid_type(::Type{<:OctahedralClenshawArray}) = OctahedralClenshawGrid +full_array_type(::Type{<:OctahedralClenshawArray}) = FullClenshawArray + +"""An `OctahedralClenshawArray` but constrained to `N=1` dimensions (horizontal only) and data is a `Vector{T}`.""" +OctahedralClenshawGrid + +# SIZE +nlat_odd(::Type{<:OctahedralClenshawArray}) = true +npoints_pole(::Type{<:OctahedralClenshawArray}) = 16 +npoints_added_per_ring(::Type{<:OctahedralClenshawArray}) = 4 + +function get_npoints2D(::Type{<:OctahedralClenshawArray}, nlat_half::Integer) + m, o = npoints_added_per_ring(OctahedralClenshawArray), npoints_pole(OctahedralClenshawArray) + return max(0, m*nlat_half^2 + 2o*nlat_half - o) # to avoid negative for nlat_half = 0 +end + +function get_nlat_half(::Type{<:OctahedralClenshawArray}, npoints2D::Integer) + m, o = npoints_added_per_ring(OctahedralClenshawArray), npoints_pole(OctahedralClenshawArray) + return round(Int, -o/m + sqrt(((o/m)^2 + (o+npoints2D)/m))) +end + +function get_nlon_per_ring(Grid::Type{<:OctahedralClenshawArray}, nlat_half::Integer, j::Integer) + nlat = get_nlat(Grid, nlat_half) + @assert 0 < j <= nlat "Ring $j is outside O$nlat_half grid." + m, o = npoints_added_per_ring(OctahedralClenshawArray), npoints_pole(OctahedralClenshawArray) + j = j > nlat_half ? nlat - j + 1 : j # flip north south due to symmetry + return o + m*j +end + +matrix_size(grid::Grid) where {Grid<:OctahedralClenshawGrid} = matrix_size(Grid, grid.nlat_half) +function matrix_size(::Type{OctahedralClenshawGrid}, nlat_half::Integer) + m, o = npoints_added_per_ring(OctahedralClenshawArray), npoints_pole(OctahedralClenshawArray) + m != 4 && @warn "This algorithm has not been generalised for m!=4." + N = (o + 4nlat_half)÷2 + return (N, N) +end + +Base.Matrix(G::OctahedralClenshawGrid{T}; kwargs...) where T = + Matrix!(zeros(T, matrix_size(G)...), G; kwargs...) + +## COORDINATES +get_colat(::Type{<:OctahedralClenshawArray}, nlat_half::Integer) = get_colat(FullClenshawArray, nlat_half) +function get_lon_per_ring(Grid::Type{<:OctahedralClenshawArray}, nlat_half::Integer, j::Integer) + nlon = get_nlon_per_ring(Grid, nlat_half, j) + return collect(0:2π/nlon:2π-π/nlon) +end + +## QUADRATURE +get_quadrature_weights(::Type{<:OctahedralClenshawArray}, nlat_half::Integer) = + clenshaw_curtis_weights(nlat_half) + +## INDEXING +function each_index_in_ring(Grid::Type{<:OctahedralClenshawArray}, + j::Integer, # ring index north to south + nlat_half::Integer) # resolution param + + nlat = get_nlat(Grid, nlat_half) + + # TODO make m, o dependent + m, o = npoints_added_per_ring(OctahedralClenshawArray), npoints_pole(OctahedralClenshawArray) + m != 4 || o != 16 && @warn "This algorithm has not been generalised for m!=4, o!=16." + + @boundscheck 0 < j <= nlat || throw(BoundsError) # ring index valid? + if j <= nlat_half # northern hemisphere incl Equator + index_1st = 2j*(j+7) - 15 # first in-ring index i + index_end = 2j*(j+9) # last in-ring index i + else # southern hemisphere excl Equator + j = nlat - j + 1 # mirror ring index around Equator + n = get_npoints2D(Grid, nlat_half) + 1 # number of grid points + 1 + index_1st = n - 2j*(j+9) # count backwards + index_end = n - (2j*(j+7) - 15) + end + return index_1st:index_end # range of i's in ring +end + +function each_index_in_ring!( rings::Vector{<:UnitRange{<:Integer}}, + Grid::Type{<:OctahedralClenshawArray}, + nlat_half::Integer) # resolution param + + nlat = length(rings) + @boundscheck nlat == get_nlat(Grid, nlat_half) || throw(BoundsError) + m, o = npoints_added_per_ring(OctahedralClenshawArray), npoints_pole(OctahedralClenshawArray) + + index_end = 0 + @inbounds for j in 1:nlat_half # North incl Eq only + index_1st = index_end + 1 # 1st index is +1 from prev ring's last index + index_end += o + m*j # add number of grid points per ring + rings[j] = index_1st:index_end # turn into UnitRange + end + @inbounds for (j, j_mirrored) in zip( nlat_half+1:nlat, # South only + nlat-nlat_half:-1:1) # reverse index + + index_1st = index_end + 1 # 1st index is +1 from prev ring's last index + index_end += o + m*j_mirrored # add number of grid points per ring + rings[j] = index_1st:index_end # turn into UnitRange + end +end + +## CONVERSION +""" +$(TYPEDSIGNATURES) +Sorts the gridpoints in `G` into the matrix `M` without interpolation. +Every quadrant of the grid `G` is rotated as specified in `quadrant_rotation`, +0 is no rotation, 1 is 90˚ clockwise, 2 is 180˚ etc. Grid quadrants are counted +eastward starting from 0˚E. The grid quadrants are moved into the matrix quadrant +(i, j) as specified. Defaults are equivalent to centered at 0˚E and a rotation +such that the North Pole is at M's midpoint.""" +Matrix!(M::AbstractMatrix, G::OctahedralClenshawGrid; kwargs...) = Matrix!((M, G); kwargs...) + +""" +$(TYPEDSIGNATURES) +Like `Matrix!(::AbstractMatrix, ::OctahedralClenshawGrid)` but for simultaneous +processing of tuples `((M1, G1), (M2, G2), ...)` with matrices `Mi` and grids `Gi`. +All matrices and grids have to be of the same size respectively.""" +function Matrix!( MGs::Tuple{AbstractMatrix{T}, OctahedralClenshawGrid}...; + quadrant_rotation::NTuple{4, Integer}=(0, 1, 2, 3), # = 0˚, 90˚, 180˚, 270˚ anti-clockwise + matrix_quadrant::NTuple{4, Tuple{Integer, Integer}}=((2, 2), (1, 2), (1, 1), (2, 1)), + ) where T + + # TODO make m, o dependent + m, o = npoints_added_per_ring(OctahedralGaussianArray), npoints_pole(OctahedralGaussianArray) + m != 4 || o != 16 && @warn "This algorithm has not been generalised for m!=4, o!=16." + + ntuples = length(MGs) + + # check that the first (matrix, grid) tuple has corresponding sizes + M, G = MGs[1] + m, n = size(M) + @boundscheck m == n || throw(BoundsError) + @boundscheck m == 2*(4+G.nlat_half) || throw(BoundsError) + + for MG in MGs # check that all matrices and all grids are of same size + Mi, Gi = MG + @boundscheck size(Mi) == size(M) || throw(BoundsError) + @boundscheck size(Gi) == size(G) || throw(BoundsError) + end + + for q in matrix_quadrant # check always in 2x2 + sr, sc = q + @boundscheck ((sr in (1, 2)) && (sc in (1, 2))) || throw(BoundsError) + end + + rings = eachring(G) # index ranges for all rings + nlat_half = G.nlat_half # number of latitude rings on one hemisphere incl Equator + nside = 4+G.nlat_half # side length of a basepixel matrix + + # sort grid indices from G into matrix M + # 1) loop over each grid point per ring + # 2) determine quadrant (0, 1, 2, 3) via modulo + # 3) get longitude index iq within quadrant + # 4) determine corresponding indices r, c in matrix M + + @inbounds for (j, ring) in enumerate(rings) + nlon = length(ring) # number of grid points in ring + for ij in ring # continuous index in grid + i = ij-ring[1] # 0-based index in ring + grid_quadrant = floor(Int, mod(4*i/nlon, 4)) # either 0, 1, 2, 3 + iq = i - grid_quadrant*(nlon÷4) # 0-based index i relative to quadrant + r = 4+min(j, nlat_half) - iq # row in matrix m (1-based) + c = (iq+1) + max(0, j-nlat_half) # column in matrix m (1-based) + + # rotate indices in quadrant + r, c = rotate_matrix_indices(r, c, nside, quadrant_rotation[grid_quadrant+1]) + + # shift grid quadrant to matrix quadrant + sr, sc = matrix_quadrant[grid_quadrant+1] + r += (sr-1)*nside # shift row into matrix quadrant + c += (sc-1)*nside # shift column into matrix quadrant + + for (Mi, Gi) in MGs # for every (matrix, grid) tuple + Mi[r, c] = convert(T, Gi[ij]) # convert data and copy over + end + end + end + + ntuples == 1 && return M + return Tuple(Mi for (Mi, Gi) in MGs) +end \ No newline at end of file diff --git a/src/RingGrids/grids/octahedral_gaussian.jl b/src/RingGrids/grids/octahedral_gaussian.jl new file mode 100644 index 000000000..b6246e2f6 --- /dev/null +++ b/src/RingGrids/grids/octahedral_gaussian.jl @@ -0,0 +1,131 @@ +"""An `OctahedralGaussianArray` is an array of octahedral grids, subtyping `AbstractReducedGridArray`, +that use Gaussian latitudes for each ring. First dimension of the underlying `N`-dimensional `data` +represents the horizontal dimension, in ring order (0 to 360˚E, then north to south), +other dimensions are used for the vertical and/or time or other dimensions. +The resolution parameter of the horizontal grid is `nlat_half` (number of latitude rings +on one hemisphere, Equator included) and the ring indices are precomputed in `rings`. + +These grids are called octahedral because after starting with 20 points on the first ring around the +north pole (default) they increase the number of longitude points for each ring by 4, such that +they can be conceptually thought of as lying on the 4 faces of an octahedron on each hemisphere. +Hence, these grids have 20, 24, 28, ... longitude points for ring 1, 2, 3, ... There is no +ring on the Equator and the two rings around it have 16 + 4nlat_half longitude points before +reducing the number of longitude points per ring by 4 towards the southern-most ring +j = nlat. `rings` are the precomputed ring indices, the the example above +`rings = [1:20, 21:44, 45:72, ...]`. For efficient looping see `eachring` and `eachgrid`. +Fields are +$(TYPEDFIELDS)""" +struct OctahedralGaussianArray{T, N, ArrayType <: AbstractArray{T, N}} <: AbstractReducedGridArray{T, N, ArrayType} + data::ArrayType # data array, ring by ring, north to south + nlat_half::Int # number of latitudes on one hemisphere + rings::Vector{UnitRange{Int}} # TODO make same array type as data? + + OctahedralGaussianArray(data::A, nlat_half, rings) where {A <: AbstractArray{T, N}} where {T, N} = + check_inputs(data, nlat_half, rings, OctahedralGaussianArray) ? + new{T, N, A}(data, nlat_half, rings) : + error_message(data, nlat_half, rings, OctahedralGaussianArray, T, N, A) +end + +# TYPES +const OctahedralGaussianGrid{T} = OctahedralGaussianArray{T, 1, Vector{T}} +nonparametric_type(::Type{<:OctahedralGaussianArray}) = OctahedralGaussianArray +horizontal_grid_type(::Type{<:OctahedralGaussianArray}) = OctahedralGaussianGrid +full_array_type(::Type{<:OctahedralGaussianArray}) = FullGaussianArray + +"""An `OctahedralGaussianArray` but constrained to `N=1` dimensions (horizontal only) and data is a `Vector{T}`.""" +OctahedralGaussianGrid + +# SIZE +nlat_odd(::Type{<:OctahedralGaussianArray}) = false + +"""$(TYPEDSIGNATURES) [EXPERIMENTAL] additional number of longitude points on the first and last ring. +Change to 0 to start with 4 points on the first ring.""" +npoints_pole(::Type{<:OctahedralGaussianArray}) = 16 + +"""$(TYPEDSIGNATURES) [EVEN MORE EXPERIMENTAL] number of longitude points added (removed) for every ring +towards the Equator (on the southern hemisphere towards the south pole).""" +npoints_added_per_ring(::Type{<:OctahedralGaussianArray}) = 4 + +function get_npoints2D(::Type{<:OctahedralGaussianArray}, nlat_half::Integer) + m, o = npoints_added_per_ring(OctahedralGaussianArray), npoints_pole(OctahedralGaussianArray) + return m*nlat_half^2 + (2o+m)*nlat_half +end + +function get_nlat_half(::Type{<:OctahedralGaussianArray}, npoints2D::Integer) + m, o = npoints_added_per_ring(OctahedralGaussianArray), npoints_pole(OctahedralGaussianArray) + return round(Int, -(2o + m)/2m + sqrt(((2o+m)/2m)^2 + npoints2D/m)) +end + +function get_nlon_per_ring(Grid::Type{<:OctahedralGaussianArray}, nlat_half::Integer, j::Integer) + nlat = get_nlat(Grid, nlat_half) + @assert 0 < j <= nlat "Ring $j is outside O$nlat_half grid." + m, o = npoints_added_per_ring(OctahedralGaussianArray), npoints_pole(OctahedralGaussianArray) + j = j > nlat_half ? nlat - j + 1 : j # flip north south due to symmetry + return o + m*j +end + +# maybe define at some point for Matrix(::OctahedralGaussianGrid) +# matrix_size(G::OctahedralGaussianGrid) = (2*(4+G.nlat_half), 2*(4+G.nlat_half+1)) + +## COORDINATES +get_colat(::Type{<:OctahedralGaussianArray}, nlat_half::Integer) = get_colat(FullGaussianArray, nlat_half) +function get_lon_per_ring(Grid::Type{<:OctahedralGaussianArray}, nlat_half::Integer, j::Integer) + nlon = get_nlon_per_ring(Grid, nlat_half, j) + return collect(0:2π/nlon:2π-π/nlon) +end + +## QUADRATURE +get_quadrature_weights(::Type{<:OctahedralGaussianArray}, nlat_half::Integer) = gaussian_weights(nlat_half) + +## INDEXING + +"""$(TYPEDSIGNATURES) `UnitRange{Int}` to index grid points on ring `j` of `Grid` +at resolution `nlat_half`.""" +function each_index_in_ring(Grid::Type{<:OctahedralGaussianArray}, + j::Integer, # ring index north to south + nlat_half::Integer) # resolution param + + nlat = get_nlat(Grid, nlat_half) + + # TODO make m, o dependent + m, o = npoints_added_per_ring(OctahedralGaussianArray), npoints_pole(OctahedralGaussianArray) + m != 4 || o != 16 && @warn "This algorithm has not been generalised for m!=4, o!=16." + + @boundscheck 0 < j <= nlat || throw(BoundsError) # ring index valid? + if j <= nlat_half # northern hemisphere incl Equator + index_1st = 2j*(j+7) - 15 # first in-ring index i + index_end = 2j*(j+9) # last in-ring index i + else # southern hemisphere excl Equator + j = nlat - j + 1 # mirror ring index around Equator + n = get_npoints2D(Grid, nlat_half) + 1 # number of grid points + 1 + index_1st = n - 2j*(j+9) # count backwards + index_end = n - (2j*(j+7) - 15) + end + return index_1st:index_end # range of i's in ring +end + +"""$(TYPEDSIGNATURES) precompute a `Vector{UnitRange{Int}} to index grid points on +every ring `j` (elements of the vector) of `Grid` at resolution `nlat_half`. +See `eachring` and `eachgrid` for efficient looping over grid points.""" +function each_index_in_ring!( rings::Vector{<:UnitRange{<:Integer}}, + Grid::Type{<:OctahedralGaussianArray}, + nlat_half::Integer) # resolution param + + nlat = length(rings) + @boundscheck nlat == get_nlat(Grid, nlat_half) || throw(BoundsError) + m, o = npoints_added_per_ring(OctahedralGaussianArray), npoints_pole(OctahedralGaussianArray) + + index_end = 0 + @inbounds for j in 1:nlat_half # North incl Eq only + index_1st = index_end + 1 # 1st index is +1 from prev ring's last index + index_end += o + m*j # add number of grid points per ring + rings[j] = index_1st:index_end # turn into UnitRange + end + @inbounds for (j, j_mirrored) in zip( nlat_half+1:nlat, # South only + nlat-nlat_half:-1:1) # reverse index + + index_1st = index_end + 1 # 1st index is +1 from prev ring's last index + index_end += o + m*j_mirrored # add number of grid points per ring + rings[j] = index_1st:index_end # turn into UnitRange + end +end \ No newline at end of file diff --git a/src/RingGrids/grids_general.jl b/src/RingGrids/grids_general.jl deleted file mode 100644 index b7dcae803..000000000 --- a/src/RingGrids/grids_general.jl +++ /dev/null @@ -1,254 +0,0 @@ -abstract type AbstractGrid{T} <: AbstractVector{T} end - -# all AbstractGrids have their grid points stored in a vector field `data` -# propagate length, size, getindex, setindex! for that -Base.length(G::AbstractGrid) = length(G.data) -Base.size(G::AbstractGrid) = size(G.data) -@inline function Base.getindex(G::AbstractGrid, k::Integer) - @boundscheck 0 < k <= length(G.data) || throw(BoundsError(G, k)) - @inbounds r = G.data[k] - return r -end - -@inline Base.setindex!(G::AbstractGrid, x, k::Integer) = setindex!(G.data, x, k) - -# with ranges -@inline Base.getindex(G::AbstractGrid, r::AbstractRange) = G.data[r] -@inline Base.setindex!(G::AbstractGrid, x::AbstractVector, r::AbstractRange) = setindex!(G.data, x, r) - -function grids_match(A::AbstractGrid, B::AbstractGrid) - length(A) == length(B) && return _grids_match(typeof(A), typeof(B)) - return false -end - -function _grids_match(A::Type{<:AbstractGrid}, B::Type{<:AbstractGrid}) - # throws an error for non-parametric types... - return A.name.wrapper == B.name.wrapper -end - -""" - abstract type AbstractGrid{T} <: AbstractVector{T} end - -The abstract supertype for all spatial grids on the sphere supported by SpeedyWeather.jl. -Every new grid has to be of the form - - abstract type AbstractGridClass{T} <: AbstractGrid{T} end - struct MyNewGrid{T} <: AbstractGridClass{T} - data::Vector{T} # all grid points unravelled into a vector - nlat_half::Int # resolution: latitude rings on one hemisphere (Equator incl) - end - -`MyNewGrid` should belong to a grid class like `AbstractFullGrid`, `AbstractOctahedralGrid` or -`AbstractHEALPixGrid` (that already exist but you may introduce a new class of grids) that share -certain features such as the number of longitude points per latitude ring and indexing, but may -have different latitudes or offset rotations. Each new grid `Grid` (or grid class) then has to -implement the following methods (as an example, see octahedral.jl) - -Fundamental grid properties - get_npoints # total number of grid points - nlat_odd # does the grid have an odd number of latitude rings? - get_nlat # total number of latitude rings - get_nlat_half # number of latitude rings on one hemisphere incl Equator - -Indexing - get_nlon_max # maximum number of longitudes points (at the Equator) - get_nlon_per_ring # number of longitudes on ring j - each_index_in_ring # a unit range that indexes all longitude points on a ring - -Coordinates - get_colat # vector of colatitudes (radians) - get_colatlon # vectors of colatitudes, longitudes (both radians) - -Spectral truncation - truncation_order # linear, quadratic, cubic = 1, 2, 3 for grid - get_truncation # spectral truncation given a grid resolution - get_resolution # grid resolution given a spectral truncation - -Quadrature weights and solid angles - get_quadrature_weights # = sinθ Δθ for grid points on ring j for meridional integration - get_solid_angle # = sinθ Δθ Δϕ, solid angle of grid points on ring j -""" -AbstractGrid - -# Define methods that are universally applicable to any G<:AbstractGrid here -# generator functions for grid -Base.zeros(::Type{Grid}, nlat_half::Integer) where {Grid<:AbstractGrid{T}} where T = - Grid(zeros(T, get_npoints(Grid, nlat_half)), nlat_half) -# use Float64 if not provided -Base.zeros(::Type{Grid}, nlat_half::Integer) where {Grid<:AbstractGrid} = zeros(Grid{Float64}, nlat_half) -# zero element of an AbstractGrid instance grid by packing a zero(::Vector) into grid -Base.zero(grid::Grid) where {Grid<:AbstractGrid} = Grid(zero(grid.data)) - -# initialise with ones -Base.ones(::Type{Grid}, nlat_half::Integer) where {Grid<:AbstractGrid{T}} where T = - Grid(ones(T, get_npoints(Grid, nlat_half)), nlat_half) -Base.ones(::Type{Grid}, nlat_half::Integer) where {Grid<:AbstractGrid} = ones(Grid{Float64}, nlat_half) - -# in case type parameter T in Grid{T} is provided -function (::Type{Grid})(::UndefInitializer, nlat_half::Integer) where {Grid<:AbstractGrid{T}} where T - return Grid(Vector{T}(undef, get_npoints(Grid, nlat_half)), nlat_half) -end - -# use Float64 if not -function (::Type{Grid})(::UndefInitializer, nlat_half::Integer) where {Grid<:AbstractGrid} - return Grid(Vector{Float64}(undef, get_npoints(Grid, nlat_half)), nlat_half) -end - -# randn initializer, use Float64 if T in Grid{T} not provided -Base.randn(::Type{Grid}, nlat_half::Integer) where {Grid<:AbstractGrid} = randn(Grid{Float64}, nlat_half) -Base.randn(::Type{Grid}, nlat_half::Integer) where {Grid<:AbstractGrid{T}} where T = - Grid(randn(T, get_npoints(Grid, nlat_half)), nlat_half) - -# rand initializer, use Float64 if T in Grid{T} not provided -Base.rand(::Type{Grid}, nlat_half::Integer) where {Grid<:AbstractGrid} = rand(Grid{Float64}, nlat_half) -Base.rand(::Type{Grid}, nlat_half::Integer) where {Grid<:AbstractGrid{T}} where T = - Grid(rand(T, get_npoints(Grid, nlat_half)), nlat_half) - -# truncation is the spectral truncation corresponding to size of grid and lin/quad/cubic truncation -get_resolution(grid::AbstractGrid) = get_nlat_half(grid) -get_nlat_half(grid::AbstractGrid) = grid.nlat_half - -# does the grid have an odd number of latitudes? -nlat_odd(grid::AbstractGrid) = nlat_odd(typeof(grid)) - -# get total number of latitude rings, *(nlat_half > 0) to return 0 for nlat_half = 0 -get_nlat(Grid::Type{<:AbstractGrid}, nlat_half::Integer) = 2nlat_half - nlat_odd(Grid)*(nlat_half > 0) -get_nlat(grid::Grid) where {Grid<:AbstractGrid} = get_nlat(Grid, grid.nlat_half) - -# get total number of grid pointst -get_npoints(grid::Grid) where {Grid<:AbstractGrid} = get_npoints(Grid, grid.nlat_half) - -# coordinates -get_latdlonds(grid::Grid) where {Grid<:AbstractGrid} = get_latdlonds(Grid, grid.nlat_half) - -function get_latdlonds(Grid::Type{<:AbstractGrid}, nlat_half::Integer) - colats, lons = get_colatlons(Grid, nlat_half) # colatitudes, longitudes in radians - latds = colats # flat copy rename before conversion - londs = lons - latds .= π/2 .- colats # colatiudes to latitudes in radians - latds .= latds .* (360/2π) # now in degrees 90˚...-90˚ - londs .*= (360/2π) - - return latds, londs -end - -function get_latlons(Grid::Type{<:AbstractGrid}, nlat_half::Integer) - colats, lons = get_colatlons(Grid, nlat_half) # colatitudes, longitudes in radians - lats = colats # flat copy rename before conversion - lats .= π/2 .- colats # colatiudes to latitudes in radians - - return lats, lons -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) -get_lon(grid::Grid) where {Grid<:AbstractGrid} = get_lon(Grid, grid.nlat_half) -get_colat(grid::Grid) where {Grid<:AbstractGrid} = get_colat(Grid, grid.nlat_half) -get_colatlons(grid::Grid) where {Grid<:AbstractGrid} = get_colatlons(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) - return get_lat(Grid, nlat_half) * (360/2π) -end - -# only defined for full grids, empty vector as fallback -get_lon(::Type{<:AbstractGrid}, nlat_half::Integer) = Float64[] -get_lond(::Type{<:AbstractGrid}, nlat_half::Integer) = Float64[] - -""" - i = each_index_in_ring(grid, j) - -UnitRange `i` to access data on grid `grid` on ring `j`.""" -function each_index_in_ring(grid::Grid, j::Integer) where {Grid<:AbstractGrid} - return each_index_in_ring(Grid, j, grid.nlat_half) -end - -""" - ijs = eachgridpoint(grid) - -UnitRange `ijs` to access each grid point on grid `grid`.""" -eachgridpoint(grid::AbstractGrid) = Base.OneTo(get_npoints(grid)) -function eachgridpoint(grid1::Grid, grids::Grid...) where {Grid<:AbstractGrid} - n = length(grid1) - Base._all_match_first(X->length(X), n, grid1, grids...) || throw(BoundsError) - return eachgridpoint(grid1) -end - -""" -$(TYPEDSIGNATURES) -Vector{UnitRange} `rings` to loop over every ring of grid `grid` -and then each grid point per ring. To be used like - - rings = eachring(grid) - for ring in rings - for ij in ring - grid[ij]""" -eachring(grid::AbstractGrid) = eachring(typeof(grid), grid.nlat_half) - -function eachring(Grid::Type{<:AbstractGrid}, nlat_half::Integer) - rings = Vector{UnitRange{Int}}(undef, get_nlat(Grid, nlat_half)) - each_index_in_ring!(rings, Grid, nlat_half) - return rings -end - -""" -$(TYPEDSIGNATURES) -Same as `eachring(grid)` but performs a bounds check to assess that all grids -in `grids` are of same size.""" -function eachring(grid1::Grid, grids::Grid...) where {Grid<:AbstractGrid} - @inline - n = length(grid1) - Base._all_match_first(X->length(X), n, grid1, grids...) || throw(BoundsError) - return eachring(grid1) -end - -""" -$(TYPEDSIGNATURES) -Obtain ring index j from gridpoint ij and Vector{UnitRange} describing rind indices -as obtained from eachring(::Grid)""" -function whichring(ij::Integer, rings::Vector{UnitRange{Int}}) - @boundscheck 0 < ij <= rings[end][end] || throw(BoundsError) - j = 1 - @inbounds while ij > rings[j][end] - j += 1 - end - return j -end - -""" -$(TYPEDSIGNATURES) -Returns a vector `nlons` for the number of longitude points per latitude ring, north to south. -Provide grid `Grid` and its resolution parameter `nlat_half`. For both_hemisphere==false only -the northern hemisphere (incl Equator) is returned.""" -function get_nlons(Grid::Type{<:AbstractGrid}, nlat_half::Integer; both_hemispheres::Bool=false) - n = both_hemispheres ? get_nlat(Grid, nlat_half) : nlat_half - return [get_nlon_per_ring(Grid, nlat_half, j) for j in 1:n] -end - -get_nlon_max(grid::Grid) where {Grid<:AbstractGrid} = get_nlon_max(Grid, grid.nlat_half) - -## BROADCASTING -# following https://docs.julialang.org/en/v1/manual/interfaces/#man-interfaces-broadcasting -import Base.Broadcast: BroadcastStyle, Broadcasted - -# {1} as grids are <:AbstractVector, Grid here is the non-parameteric Grid type! -struct AbstractGridStyle{Grid} <: Broadcast.AbstractArrayStyle{1} end - -# important to remove Grid{T} parameter T (eltype/number format) here to broadcast -# automatically across the same grid type but with different T -# e.g. FullGaussianGrid{Float32} and FullGaussianGrid{Float64} -Base.BroadcastStyle(::Type{Grid}) where {Grid<:AbstractGrid} = - AbstractGridStyle{nonparametric_type(Grid)}() - -# allocation for broadcasting, create a new Grid with undef of type/number format T -function Base.similar(bc::Broadcasted{AbstractGridStyle{Grid}}, ::Type{T}) where {Grid, T} - Grid(Vector{T}(undef,length(bc))) -end - -# ::Val{0} for broadcasting with 0-dimensional, ::Val{1} for broadcasting with vectors -AbstractGridStyle{Grid}(::Val{0}) where Grid = AbstractGridStyle{Grid}() -AbstractGridStyle{Grid}(::Val{1}) where Grid = AbstractGridStyle{Grid}() \ No newline at end of file diff --git a/src/RingGrids/healpix.jl b/src/RingGrids/healpix.jl deleted file mode 100644 index 38846ada1..000000000 --- a/src/RingGrids/healpix.jl +++ /dev/null @@ -1,150 +0,0 @@ -""" - abstract type AbstractHEALPixGrid{T} <: AbstractGrid{T} end - -An `AbstractHEALPixGrid` is a horizontal grid similar to the standard HEALPixGrid, -but different latitudes can be used, the default HEALPix latitudes or others.""" -abstract type AbstractHEALPixGrid{T} <: AbstractGrid{T} end - -npoints_healpix(nlat_half::Integer) = 3*nlat_half^2 -nside_healpix(nlat_half::Integer) = nlat_half÷2 -nlat_half_healpix(npoints::Integer) = round(Int, sqrt(npoints/3)) # inverse of npoints_healpix -nlon_healpix(nlat_half::Integer, j::Integer) = min(4j, 2nlat_half, 8nlat_half-4j) -nlon_max_healpix(nlat_half::Integer) = 2nlat_half - -nlat_odd(::Type{<:AbstractHEALPixGrid}) = true -get_nlon_max(::Type{<:AbstractHEALPixGrid}, nlat_half::Integer) = nlon_max_healpix(nlat_half) - -function get_nlon_per_ring(G::Type{<:AbstractHEALPixGrid}, nlat_half::Integer, j::Integer) - nlat = get_nlat(G, nlat_half) - @assert 0 < j <= nlat "Ring $j is outside H$nlat_half grid." - return nlon_healpix(nlat_half, j) -end - -get_npoints(::Type{<:AbstractHEALPixGrid}, nlat_half::Integer) = npoints_healpix(nlat_half) - -function get_colatlons(Grid::Type{<:AbstractHEALPixGrid}, nlat_half::Integer) - nlat = get_nlat(Grid, nlat_half) - npoints = get_npoints(Grid, nlat_half) - nside = nside_healpix(nlat_half) - colat = get_colat(Grid, nlat_half) - - colats = zeros(npoints) - lons = zeros(npoints) - - ij = 1 - for j in 1:nlat - nlon = get_nlon_per_ring(Grid, nlat_half, j) - - # s = 1 for polar caps, s=2, 1, 2, 1, ... in the equatorial zone - s = (j < nside) || (j >= 3nside) ? 1 : ((j - nside) % 2 + 1) - lon = [π/(nlon÷2)*(i - s/2) for i in 1:nlon] - - colats[ij:ij+nlon-1] .= colat[j] - lons[ij:ij+nlon-1] .= lon - - ij += nlon - end - - return colats, lons -end - -function each_index_in_ring(::Type{<:AbstractHEALPixGrid}, # function for HEALPix grids - j::Integer, # ring index north to south - nlat_half::Integer) # resolution param - nside = nside_healpix(nlat_half) - - @boundscheck 0 < j < 4nside || throw(BoundsError) # ring index valid? - if j < nside # northern polar cap - index_1st = 2j*(j-1) + 1 # first in-ring index i - index_end = 2j*(j+1) # last in-ring index i - elseif j <= 3nside # equatorial zone with const nlon - n = 2nside^2-2nside # points in polar cap - nlon = 4nside # points on latitude rings - j = j - nside + 1 # offset ring index into eq zone - index_1st = n + (j-1)*nlon + 1 # add constant nlon per ring - index_end = n + j*nlon - else # southern polar cap - n = 12nside^2 # total number of points - j = 4nside - j # count ring index from south pole - index_1st = n - 2j*(j+1) + 1 # count backwards - index_end = n - 2j*(j-1) - end - return index_1st:index_end # range of i's in ring -end - -function each_index_in_ring!( rings::Vector{<:UnitRange{<:Integer}}, - Grid::Type{<:AbstractHEALPixGrid}, - nlat_half::Integer) # resolution param - - nlat = length(rings) - @boundscheck nlat == get_nlat(Grid, nlat_half) || throw(BoundsError) - - index_end = 0 - nside = nside_healpix(nlat_half) # side length of a basepixel - - # North polar cap - @inbounds for j in 1:nside-1 - index_1st = index_end + 1 # 1st index is +1 from prev ring's last index - index_end += 4j # add number of grid points per ring - rings[j] = index_1st:index_end # turn into UnitRange - end - - # Equatorial belt - nlon_max = get_nlon_max(Grid, nlat_half) # number of grid points on belt - @inbounds for j in nside:3nside - index_1st = index_end + 1 # 1st index is +1 from prev ring's last index - index_end += nlon_max # nlon constant in belt - rings[j] = index_1st:index_end # turn into UnitRange - end - - # South polar cap - @inbounds for (j, j_mir) in zip( 3nside+1:nlat, # South only - nside-1:-1:1) # mirror index - - index_1st = index_end + 1 # 1st index is +1 from prev ring's last index - index_end += 4j_mir # add number of grid points per ring - rings[j] = index_1st:index_end # turn into UnitRange - end -end - -""" - H = HEALPixGrid{T} - -A HEALPix grid with 12 faces, each `nside`x`nside` grid points, each covering the same area. -The number of latitude rings on one hemisphere (incl Equator) `nlat_half` is used as resolution parameter. -The values of all grid points are stored in a vector field `v` that unravels the data 0 to 360˚, -then ring by ring, which are sorted north to south.""" -struct HEALPixGrid{T} <: AbstractHEALPixGrid{T} - data::Vector{T} # data vector, ring by ring, north to south - nlat_half::Int # number of latitude rings on one hemisphere (Equator included) - - HEALPixGrid{T}(data::AbstractVector, nlat_half::Integer) where T = length(data) == npoints_healpix(nlat_half) ? - new(data, nlat_half) : error("$(length(data))-element Vector{$(eltype(data))}"* - "cannot be used to create an H$nlat_half HEALPixGrid{$T}.") -end - -nonparametric_type(::Type{<:HEALPixGrid}) = HEALPixGrid - -# infer nlat_half from data vector length, infer parametric type from eltype of data -HEALPixGrid{T}(data::AbstractVector) where T = HEALPixGrid{T}(data, nlat_half_healpix(length(data))) -HEALPixGrid(data::AbstractVector, n::Integer...) = HEALPixGrid{eltype(data)}(data, n...) - -function get_colat(::Type{<:HEALPixGrid}, nlat_half::Integer) - nlat_half == 0 && return Float64[] - - nlat = get_nlat(HEALPixGrid, nlat_half) - nside = nside_healpix(nlat_half) - colat = zeros(nlat) - - for j in 1:nside colat[j] = acos(1-j^2/3nside^2) end # north polar cap - for j in nside+1:3nside colat[j] = acos(4/3 - 2j/3nside) end # equatorial belt - for j in 3nside+1:nlat colat[j] = acos((2nlat_half-j)^2/3nside^2-1) end # south polar cap - - return colat -end - -full_grid(::Type{<:HEALPixGrid}) = FullHEALPixGrid # the full grid with same latitudes -function matrix_size(G::HEALPixGrid) - nside = nside_healpix(G.nlat_half) - return (5nside, 5nside) -end \ No newline at end of file diff --git a/src/RingGrids/interpolation.jl b/src/RingGrids/interpolation.jl index 3ca9fa0fb..06480b492 100644 --- a/src/RingGrids/interpolation.jl +++ b/src/RingGrids/interpolation.jl @@ -14,7 +14,7 @@ struct GridGeometry{G<:AbstractGrid} lon_offsets::Vector{Float64} # longitude offsets of first grid point per ring end -GridGeometry(grid::AbstractGrid) = GridGeometry(typeof(grid), grid.nlat_half) +GridGeometry(grid::AbstractGridArray) = GridGeometry(horizontal_grid_type(grid), grid.nlat_half) """ G = GridGeometry( Grid::Type{<:AbstractGrid}, @@ -26,14 +26,14 @@ unravelled indices ij.""" function GridGeometry( Grid::Type{<:AbstractGrid}, # which grid to calculate the geometry for nlat_half::Integer) # resolution parameter number of rings - nlat = get_nlat(Grid, nlat_half) # total number of latitude rings - npoints = get_npoints(Grid, nlat_half) # total number of grid points + nlat = get_nlat(Grid, nlat_half) # total number of latitude rings + npoints = get_npoints(Grid, nlat_half) # total number of grid points # LATITUDES - colat = get_colat(Grid, nlat_half) # colatitude in radians + colat = get_colat(Grid, nlat_half) # colatitude in radians lat = π/2 .- colat # latitude in radians latd = lat*360/2π # 90˚...-90˚, in degrees - latd_poles = cat(90, latd, -90, dims=1) # latd, but poles incl + latd_poles = cat(90, latd, -90, dims=1) # latd, but poles incl # Hack: use -90.00...1˚N instead of exactly -90˚N for the <=, > comparison # in find_rings! that way the last ring to the south pole can be an open @@ -42,12 +42,11 @@ function GridGeometry( Grid::Type{<:AbstractGrid}, # which grid to calculate th latd_poles[end] = latd_poles[end] - eps(latd_poles[end]) # COORDINATES for every grid point in ring order - _, londs = get_latdlonds(Grid, nlat_half) # in degrees [0˚...360˚E] + _, londs = get_latdlonds(Grid, nlat_half) # in degrees [0˚...360˚E] # RINGS and LONGITUDE OFFSETS - rings = eachring(Grid, nlat_half) # Vector{UnitRange} descr start/end index on ring - nlons = get_nlons(Grid, nlat_half, # number of longitude points per ring - both_hemispheres=true) + rings = eachring(Grid, nlat_half) # Vector{UnitRange} descr start/end index on ring + nlons = get_nlons(Grid, nlat_half) # number of longitude points per ring, pole to pole lon_offsets = [londs[ring[1]] for ring in rings]# offset of the first point from 0˚E return GridGeometry{Grid}(nlat_half, nlat, npoints, latd_poles, londs, rings, nlons, lon_offsets) @@ -213,9 +212,9 @@ function interpolate( A::AbstractGrid{NF}, # field to interpolate I::AbstractInterpolator # indices in I are assumed to be calculated already! ) where NF # use number format from input data also for output - (; npoints ) = I.locator # number of points to interpolate onto - Aout = Vector{NF}(undef, npoints) # preallocate: onto θs, λs interpolated values of A - interpolate!(Aout, A, I) # perform interpolation, store in As + (; npoints ) = I.locator # number of points to interpolate onto + Aout = Vector{NF}(undef, npoints) # preallocate: onto θs, λs interpolated values of A + interpolate!(Aout, A, I) # perform interpolation, store in As end function interpolate!( Aout::Vector, # Out: interpolated values @@ -315,7 +314,7 @@ end function find_rings!( js::Vector{<:Integer}, # Out: ring indices j Δys::Vector, # Out: distance fractions to ring further south - θs::Vector, # latitudes to interpolate onto + θs::Vector, # latitudes to interpolate onto latd::Vector; # latitudes of the rings on the original grid unsafe::Bool=false) # skip safety checks when true @@ -511,7 +510,7 @@ function grid_cell_average!( lon_in = get_lon(input) nlon_in = length(lon_in) - # output grid coordinates, append -π, 2π to have grid points + # output grid coordinates, append -π, 2π to have grid points # towards the poles definitely included colat_out = vcat(-π, get_colat(output), 2π) _, lons_out = get_colatlons(output) @@ -522,7 +521,7 @@ function grid_cell_average!( Δϕ = 2π/length(ring) # longitude spacing on this ring # indices for lat_out are shifted as north and south pole are included - θ0 = (colat_out[j] + colat_out[j+1])/2 # northern edge + θ0 = (colat_out[j] + colat_out[j+1])/2 # northern edge θ1 = (colat_out[j+1] + colat_out[j+2])/2 # southern edge # matrix indices for input grid that lie in output grid cell diff --git a/src/RingGrids/octahedral.jl b/src/RingGrids/octahedral.jl deleted file mode 100644 index 40149e177..000000000 --- a/src/RingGrids/octahedral.jl +++ /dev/null @@ -1,238 +0,0 @@ -""" - abstract type AbstractOctahedralGrid{T} <: AbstractGrid{T} end - -An `AbstractOctahedralGrid` is a horizontal grid with 16+4i longitude -points on the latitude ring i starting with i=1 around the pole. -Different latitudes can be used, Gaussian latitudes, equi-angle latitdes, or others.""" -abstract type AbstractOctahedralGrid{T} <: AbstractGrid{T} end - -get_nlon_max(::Type{<:AbstractOctahedralGrid}, nlat_half::Integer) = nlon_octahedral(nlat_half) - -function get_nlon_per_ring(Grid::Type{<:AbstractOctahedralGrid}, nlat_half::Integer, j::Integer) - nlat = get_nlat(Grid, nlat_half) - @assert 0 < j <= nlat "Ring $j is outside O$nlat_half grid." - j = j > nlat_half ? nlat - j + 1 : j # flip north south due to symmetry - return nlon_octahedral(j) -end - -function get_colatlons(Grid::Type{<:AbstractOctahedralGrid}, nlat_half::Integer) - - colat = get_colat(Grid, nlat_half) - nlat = get_nlat(Grid, nlat_half) - - npoints = get_npoints(Grid, nlat_half) - colats = zeros(npoints) # preallocate arrays - lons = zeros(npoints) - - ij = 1 # continuous index - for j in 1:nlat # populate arrays ring by ring - nlon = get_nlon_per_ring(Grid, nlat_half, j) - lon = collect(0:2π/nlon:2π-π/nlon) - - colats[ij:ij+nlon-1] .= colat[j] - lons[ij:ij+nlon-1] .= lon - - ij += nlon - end - - return colats, lons -end - -function each_index_in_ring(Grid::Type{<:AbstractOctahedralGrid}, - j::Integer, # ring index north to south - nlat_half::Integer) # resolution param - - nlat = get_nlat(Grid, nlat_half) - @boundscheck 0 < j <= nlat || throw(BoundsError) # ring index valid? - if j <= nlat_half # northern hemisphere incl Equator - index_1st = 2j*(j+7) - 15 # first in-ring index i - index_end = 2j*(j+9) # last in-ring index i - else # southern hemisphere excl Equator - j = nlat - j + 1 # mirror ring index around Equator - n = get_npoints(Grid, nlat_half) + 1 # number of grid points + 1 - index_1st = n - 2j*(j+9) # count backwards - index_end = n - (2j*(j+7) - 15) - end - return index_1st:index_end # range of i's in ring -end - -function each_index_in_ring!( rings::Vector{<:UnitRange{<:Integer}}, - Grid::Type{<:AbstractOctahedralGrid}, - nlat_half::Integer) # resolution param - - nlat = length(rings) - @boundscheck nlat == get_nlat(Grid, nlat_half) || throw(BoundsError) - - index_end = 0 - @inbounds for j in 1:nlat_half # North incl Eq only - index_1st = index_end + 1 # 1st index is +1 from prev ring's last index - index_end += 16 + 4j # add number of grid points per ring - rings[j] = index_1st:index_end # turn into UnitRange - end - @inbounds for (j, j_mir) in zip( nlat_half+1:nlat, # South only - nlat-nlat_half:-1:1) # reverse index - - index_1st = index_end + 1 # 1st index is +1 from prev ring's last index - index_end += 16 + 4j_mir # add number of grid points per ring - rings[j] = index_1st:index_end # turn into UnitRange - end -end - - -""" - G = OctahedralGaussianGrid{T} - -An Octahedral Gaussian grid that uses `nlat` Gaussian latitudes, but a decreasing number of longitude -points per latitude ring towards the poles. Starting with 20 equi-spaced longitude points (starting at 0˚E) -on the rings around the poles, each latitude ring towards the equator has consecuitively 4 more points, -one for each face of the octahedron. E.g. 20, 24, 28, 32, ...nlon-4, nlon, nlon, nlon-4, ..., 32, 28, 24, 20. -The maximum number of longitue points is `nlon`. The values of all grid points are stored in a vector -field `v` that unravels the data 0 to 360˚, then ring by ring, which are sorted north to south.""" -struct OctahedralGaussianGrid{T} <: AbstractOctahedralGrid{T} - data::Vector{T} # data vector, ring by ring, north to south - nlat_half::Int # number of latitudes on one hemisphere - - # check that `nlat_half` match the vector `v` length - OctahedralGaussianGrid{T}(data::AbstractVector, nlat_half::Integer) where T = length(data) == npoints_octahedral(nlat_half, false) ? - new(data, nlat_half) : error("$(length(data))-element Vector{$(eltype(data))}"* - "cannot be used to create a O$(nlat_half) OctahedralGaussianGrid{$T}.") -end - -nonparametric_type(::Type{<:OctahedralGaussianGrid}) = OctahedralGaussianGrid - -# number of points and longitudes per ring on the octahedral grid -npoints_octahedral(nlat_half::Integer, nlat_oddp::Bool) = - nlat_oddp ? max(0, 4nlat_half^2 + 32nlat_half - 16) : 4nlat_half^2 + 36nlat_half # max(0, ...) needed to avoid negative array size when nlat_half==0 -nlat_half_octahedral(npoints::Integer, nlat_oddp::Bool) = - nlat_oddp ? round(Int, -4+sqrt(20 + npoints/4)) : round(Int, -9/2+sqrt((9/2)^2 + npoints/4)) # inverse -nlon_octahedral(j::Integer) = 16+4j - -# infer nside from data vector length, infer parametric type from eltype of data -OctahedralGaussianGrid{T}(data::AbstractVector) where T = OctahedralGaussianGrid{T}(data, nlat_half_octahedral(length(data), false)) -OctahedralGaussianGrid(data::AbstractVector, n::Integer...) = OctahedralGaussianGrid{eltype(data)}(data, n...) - -nlat_odd(::Type{<:OctahedralGaussianGrid}) = false -get_npoints(::Type{<:OctahedralGaussianGrid}, nlat_half::Integer) = npoints_octahedral(nlat_half, false) -get_colat(::Type{<:OctahedralGaussianGrid}, nlat_half::Integer) = get_colat(FullGaussianGrid, nlat_half) -get_quadrature_weights(::Type{<:OctahedralGaussianGrid}, nlat_half::Integer) = gaussian_weights(nlat_half) -full_grid(::Type{<:OctahedralGaussianGrid}) = FullGaussianGrid # the full grid with same latitudes -matrix_size(G::OctahedralGaussianGrid) = (2*(4+G.nlat_half), 2*(4+G.nlat_half+1)) - -""" - G = OctahedralClenshawGrid{T} - -An Octahedral Clenshaw grid that uses `nlat` equi-spaced latitudes. Like FullClenshawGrid, the central -latitude ring is on the Equator. Like OctahedralGaussianGrid, the number of longitude points per -latitude ring decreases towards the poles. Starting with 20 equi-spaced longitude points (starting at 0˚E) -on the rings around the poles, each latitude ring towards the equator has consecuitively 4 more points, -one for each face of the octahedron. E.g. 20, 24, 28, 32, ...nlon-4, nlon, nlon, nlon-4, ..., 32, 28, 24, 20. -The maximum number of longitue points is `nlon`. The values of all grid points are stored in a vector -field `v` that unravels the data 0 to 360˚, then ring by ring, which are sorted north to south.""" -struct OctahedralClenshawGrid{T} <: AbstractOctahedralGrid{T} - data::Vector{T} # data vector, ring by ring, north to south - nlat_half::Int # number of latitudes on one hemisphere (incl Equator) - - # check that `nlat_half` match the vector `v` length - OctahedralClenshawGrid{T}(data::AbstractVector, nlat_half::Integer) where T = length(data) == npoints_octahedral(nlat_half, true) ? - new(data, nlat_half) : error("$(length(data))-element Vector{$(eltype(data))}"* - "cannot be used to create a O$(nlat_half) OctahedralClenshawGrid{$T}.") -end - -nonparametric_type(::Type{<:OctahedralClenshawGrid}) = OctahedralClenshawGrid - -# infer nlat_half from data vector length, infer parametric type from eltype of data -OctahedralClenshawGrid{T}(data::AbstractVector) where T = OctahedralClenshawGrid{T}(data, - nlat_half_octahedral(length(data), true)) -OctahedralClenshawGrid(data::AbstractVector, n::Integer...) = OctahedralClenshawGrid{eltype(data)}(data, n...) - -nlat_odd(::Type{<:OctahedralClenshawGrid}) = true -get_npoints(::Type{<:OctahedralClenshawGrid}, nlat_half::Integer) = npoints_octahedral(nlat_half, true) -get_colat(::Type{<:OctahedralClenshawGrid}, nlat_half::Integer) = get_colat(FullClenshawGrid, nlat_half) -get_quadrature_weights(::Type{<:OctahedralClenshawGrid}, nlat_half::Integer) = clenshaw_curtis_weights(nlat_half) -full_grid(::Type{<:OctahedralClenshawGrid}) = FullClenshawGrid # the full grid with same latitudes - -matrix_size(G::OctahedralClenshawGrid) = (2*(4+G.nlat_half), 2*(4+G.nlat_half)) -matrix_size(::Type{OctahedralClenshawGrid}, nlat_half::Integer) = (2*(4+nlat_half), 2*(4+nlat_half)) -Base.Matrix(G::OctahedralClenshawGrid{T}; kwargs...) where T = Matrix!(zeros(T, matrix_size(G)...), G; kwargs...) - -""" - Matrix!(M::AbstractMatrix, - G::OctahedralClenshawGrid; - quadrant_rotation=(0, 1, 2, 3), - matrix_quadrant=((2, 2), (1, 2), (1, 1), (2, 1)), - ) - -Sorts the gridpoints in `G` into the matrix `M` without interpolation. -Every quadrant of the grid `G` is rotated as specified in `quadrant_rotation`, -0 is no rotation, 1 is 90˚ clockwise, 2 is 180˚ etc. Grid quadrants are counted -eastward starting from 0˚E. The grid quadrants are moved into the matrix quadrant -(i, j) as specified. Defaults are equivalent to centered at 0˚E and a rotation -such that the North Pole is at M's midpoint.""" -Matrix!(M::AbstractMatrix, G::OctahedralClenshawGrid; kwargs...) = Matrix!((M, G); kwargs...) - -""" - Matrix!(MGs::Tuple{AbstractMatrix{T}, OctahedralClenshawGrid}...; kwargs...) - -Like `Matrix!(::AbstractMatrix, ::OctahedralClenshawGrid)` but for simultaneous -processing of tuples `((M1, G1), (M2, G2), ...)` with matrices `Mi` and grids `Gi`. -All matrices and grids have to be of the same size respectively.""" -function Matrix!( MGs::Tuple{AbstractMatrix{T}, OctahedralClenshawGrid}...; - quadrant_rotation::NTuple{4, Integer}=(0, 1, 2, 3), # = 0˚, 90˚, 180˚, 270˚ anti-clockwise - matrix_quadrant::NTuple{4, Tuple{Integer, Integer}}=((2, 2), (1, 2), (1, 1), (2, 1)), - ) where T - - ntuples = length(MGs) - - # check that the first (matrix, grid) tuple has corresponding sizes - M, G = MGs[1] - m, n = size(M) - @boundscheck m == n || throw(BoundsError) - @boundscheck m == 2*(4+G.nlat_half) || throw(BoundsError) - - for MG in MGs # check that all matrices and all grids are of same size - Mi, Gi = MG - @boundscheck size(Mi) == size(M) || throw(BoundsError) - @boundscheck size(Gi) == size(G) || throw(BoundsError) - end - - for q in matrix_quadrant # check always in 2x2 - sr, sc = q - @boundscheck ((sr in (1, 2)) && (sc in (1, 2))) || throw(BoundsError) - end - - rings = eachring(G) # index ranges for all rings - nlat_half = G.nlat_half # number of latitude rings on one hemisphere incl Equator - nside = 4+G.nlat_half # side length of a basepixel matrix - - # sort grid indices from G into matrix M - # 1) loop over each grid point per ring - # 2) determine quadrant (0, 1, 2, 3) via modulo - # 3) get longitude index iq within quadrant - # 4) determine corresponding indices r, c in matrix M - - @inbounds for (j, ring) in enumerate(rings) - nlon = length(ring) # number of grid points in ring - for ij in ring # continuous index in grid - i = ij-ring[1] # 0-based index in ring - grid_quadrant = floor(Int, mod(4*i/nlon, 4)) # either 0, 1, 2, 3 - iq = i - grid_quadrant*(nlon÷4) # 0-based index i relative to quadrant - r = 4+min(j, nlat_half) - iq # row in matrix m (1-based) - c = (iq+1) + max(0, j-nlat_half) # column in matrix m (1-based) - - # rotate indices in quadrant - r, c = rotate_matrix_indices(r, c, nside, quadrant_rotation[grid_quadrant+1]) - - # shift grid quadrant to matrix quadrant - sr, sc = matrix_quadrant[grid_quadrant+1] - r += (sr-1)*nside # shift row into matrix quadrant - c += (sc-1)*nside # shift column into matrix quadrant - - for (Mi, Gi) in MGs # for every (matrix, grid) tuple - Mi[r, c] = convert(T, Gi[ij]) # convert data and copy over - end - end - end - - ntuples == 1 && return M - return Tuple(Mi for (Mi, Gi) in MGs) # -end \ No newline at end of file diff --git a/src/RingGrids/quadrature_weights.jl b/src/RingGrids/quadrature_weights.jl index 3a51673dc..dc997557c 100644 --- a/src/RingGrids/quadrature_weights.jl +++ b/src/RingGrids/quadrature_weights.jl @@ -1,21 +1,51 @@ -# QUADRATURE WEIGHTS (EXACT) -# gaussian_weights are exact for Gaussian latitudes when nlat > (2T+1)/2 -# clenshaw_curtis_weights are exact for equi-angle latitudes when nlat > 2T+1 -gaussian_weights(nlat_half::Integer) = FastGaussQuadrature.gausslegendre(2nlat_half)[2][1:nlat_half] +"""$(TYPEDSIGNATURES) +The Gaussian weights for a Gaussian grid (full or octahedral) of size nlat_half. +Gaussian weights are of length `nlat`, i.e. a vector for every latitude ring, pole to pole. +`sum(gaussian_weights(nlat_half))` is always `2` as int_0^π sin(x) dx = 2 (colatitudes), +or equivalently int_-pi/2^pi/2 cos(x) dx (latitudes). + +Integration (and therefore the spectral transform) is _exact_ (only rounding errors) +when using Gaussian grids provided that nlat >= 3(T + 1)/2, meaning that a grid resolution +of at least 96x48 (nlon x nlat) is sufficient for an exact transform with a T=31 spectral +truncation.""" +gaussian_weights(nlat_half::Integer) = FastGaussQuadrature.gausslegendre(2nlat_half)[2] + + +"""$(TYPEDSIGNATURES) +The Clenshaw-Curtis weights for a Clenshaw grid (full or octahedral) of size nlat_half. +Clenshaw-Curtis weights are of length `nlat`, i.e. a vector for every latitude ring, pole to pole. +`sum(clenshaw_curtis_weights(nlat_half))` is always `2` as int_0^π sin(x) dx = 2 (colatitudes), +or equivalently int_-pi/2^pi/2 cos(x) dx (latitudes). + +Integration (and therefore the spectral transform) is _exact_ (only rounding errors) +when using Clenshaw grids provided that nlat >= 2(T + 1), meaning that a grid resolution +of at least 128x64 (nlon x nlat) is sufficient for an exact transform with a T=31 spectral +truncation.""" function clenshaw_curtis_weights(nlat_half::Integer) - nlat = 2nlat_half - 1 + nlat = get_nlat(FullClenshawArray, nlat_half) θs = get_colat(FullClenshawGrid, nlat_half) - return [4sin(θj)/(nlat+1)*sum([sin(p*θj)/p for p in 1:2:nlat]) for θj in θs[1:nlat_half]] + return [4sin(θj)/(nlat+1)*sum([sin(p*θj)/p for p in 1:2:nlat]) for θj in θs] end -# QUADRATURE WEIGHTS (INEXACT), for HEALPix full grids -healpix_weights(nlat_half::Integer) = - [nlon_healpix(nlat_half, j) for j in 1:nlat_half].*(2/npoints_healpix(nlat_half)) -octahealpix_weights(nlat_half::Integer) = - [nlon_octahealpix(nlat_half, j) for j in 1:nlat_half].*(2/npoints_octahealpix(nlat_half)) +"""$(TYPEDSIGNATURES) +The equal-area weights used for the HEALPix grids (original or OctaHEALPix) of size nlat_half. +The weights are of length `nlat`, i.e. a vector for every latitude ring, pole to pole. +`sum(equal_area_weights(nlat_half))` is always `2` as int_0^π sin(x) dx = 2 (colatitudes), +or equivalently int_-pi/2^pi/2 cos(x) dx (latitudes). Integration (and therefore the +spectral transform) is not exact with these grids but errors reduce for higher resolution.""" +function equal_area_weights(Grid::Type{<:AbstractGridArray}, nlat_half::Integer) + nlat = get_nlat(Grid, nlat_half) + npoints2D = get_npoints2D(Grid, nlat_half) + weights = zeros(nlat) + for j in 1:nlat + nlon = get_nlon_per_ring(Grid, nlat_half, j) + weights[j] = 2nlon/npoints2D + end + return weights +end # SOLID ANGLES ΔΩ = sinθ Δθ Δϕ -get_solid_angles(Grid::Type{<:AbstractGrid}, nlat_half::Integer) = +get_solid_angles(Grid::Type{<:AbstractGridArray}, nlat_half::Integer) = get_quadrature_weights(Grid, nlat_half) .* (2π./get_nlons(Grid, nlat_half)) -get_solid_angles(Grid::Type{<:Union{HEALPixGrid, OctaHEALPixGrid}}, nlat_half::Integer) = - 4π/get_npoints(Grid, nlat_half)*ones(nlat_half) +get_solid_angles(Grid::Type{<:Union{HEALPixArray, OctaHEALPixArray}}, nlat_half::Integer) = + 4π/get_npoints2D(Grid, nlat_half)*ones(get_nlat(Grid, nlat_half)) \ No newline at end of file diff --git a/src/RingGrids/reduced_grids.jl b/src/RingGrids/reduced_grids.jl new file mode 100644 index 000000000..37985935a --- /dev/null +++ b/src/RingGrids/reduced_grids.jl @@ -0,0 +1,40 @@ +"""Subtype of `AbstractGridArray` for arrays of rings grids that have a reduced number +of longitude points towards the poles, i.e. they are not "full", see `AbstractFullGridArray`. +Data on these grids cannot be represented as matrix and has to be unravelled into a vector, +ordered 0 to 360˚E then north to south, ring by ring. Examples for reduced grids are +the octahedral Gaussian or Clenshaw grids, or the HEALPix grid.""" +abstract type AbstractReducedGridArray{T, N, ArrayType <: AbstractArray{T, N}} <: AbstractGridArray{T, N, ArrayType} end + +"""Horizontal abstract type for all `AbstractReducedGridArray` with `N=1` (i.e. horizontal only) +and `ArrayType` of `Vector{T}` with element type `T`.""" +const AbstractReducedGrid{T} = AbstractReducedGridArray{T, 1, Vector{T}} + +# only defined for full grids +get_lon(::Type{<:AbstractReducedGridArray}, nlat_half::Integer) = Float64[] +get_lond(::Type{<:AbstractReducedGridArray}, nlat_half::Integer) = Float64[] + +# all reduced grids have their maximum number of longitude points around the equator, i.e. j = nlat_half +get_nlon_max(Grid::Type{<:AbstractReducedGridArray}, nlat_half::Integer) = get_nlon_per_ring(Grid, nlat_half, nlat_half) + +function get_colatlons(Grid::Type{<:AbstractReducedGridArray}, nlat_half::Integer) + + colat = get_colat(Grid, nlat_half) + nlat = get_nlat(Grid, nlat_half) + + npoints = get_npoints(Grid, nlat_half) + colats = zeros(npoints) # preallocate arrays + lons = zeros(npoints) + + ij = 1 # continuous index + for j in 1:nlat # populate arrays ring by ring + lon = get_lon_per_ring(Grid, nlat_half, j) + nlon = length(lon) + + colats[ij:ij+nlon-1] .= colat[j] + lons[ij:ij+nlon-1] .= lon + + ij += nlon + end + + return colats, lons +end \ No newline at end of file diff --git a/src/RingGrids/scaling.jl b/src/RingGrids/scaling.jl index 36e4cc1cb..ba6f63a8a 100644 --- a/src/RingGrids/scaling.jl +++ b/src/RingGrids/scaling.jl @@ -1,29 +1,36 @@ -# 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) +# functions to scale the latitude of any grid, in-place +scale_coslat!( grid::AbstractGridArray) = _scale_coslat!(grid, power=1) +scale_coslat²!( grid::AbstractGridArray) = _scale_coslat!(grid, power=2) +scale_coslat⁻¹!(grid::AbstractGridArray) = _scale_coslat!(grid, power=-1) +scale_coslat⁻²!(grid::AbstractGridArray) = _scale_coslat!(grid, 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) +# and via a deepcopy +scale_coslat( grid::AbstractGridArray) = scale_coslat!( deepcopy(grid)) +scale_coslat²( grid::AbstractGridArray) = scale_coslat²!( deepcopy(grid)) +scale_coslat⁻¹(grid::AbstractGridArray) = scale_coslat⁻¹!(deepcopy(grid)) +scale_coslat⁻²(grid::AbstractGridArray) = scale_coslat⁻²!(deepcopy(grid)) + +function _scale_coslat!(grid::Grid; power=1) where {Grid<:AbstractGridArray} + lat = get_lat(Grid, grid.nlat_half) + T = eltype(grid) + coslat = @. convert(T, cos(lat)^power) + return _scale_lat!(grid, 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) +function _scale_lat!(grid::AbstractGridArray{T}, v::AbstractVector) where T + @boundscheck get_nlat(grid) == length(v) || throw(BoundsError) - @inbounds for (j, ring) in enumerate(rings) - vj = convert(NF, v[j]) - for ij in ring - A[ij] *= vj + @inbounds for k in eachgrid(grid) + for (j, ring) in enumerate(eachring(grid)) + vj = convert(T, v[j]) + for ij in ring + grid[ij, k] *= vj + end end end - return A -end \ No newline at end of file + return grid +end \ No newline at end of file diff --git a/src/RingGrids/show.jl b/src/RingGrids/show.jl index 88304d088..51ea94095 100644 --- a/src/RingGrids/show.jl +++ b/src/RingGrids/show.jl @@ -1,13 +1,13 @@ # For grids, also add info about the number of rings, e.g. # julia> A = randn(OctahedralClenshawGrid, 24) # 3056-element, 47-ring OctahedralClenshawGrid{Float64}: -function Base.array_summary(io::IO, grid::AbstractGrid, inds::Tuple{Vararg{Base.OneTo}}) +function Base.array_summary(io::IO, grid::AbstractGridArray, inds::Tuple{Vararg{Base.OneTo}}) print(io, Base.dims2string(length.(inds)), ", $(get_nlat(grid))-ring ") Base.showarg(io, grid, true) end function plot(A::AbstractGrid; title::String="$(get_nlat(A))-ring $(typeof(A))") - A_full = interpolate(full_grid(typeof(A)), A.nlat_half, A) + A_full = interpolate(full_grid_type(A), A.nlat_half, A) plot(A_full; title) end diff --git a/src/RingGrids/similar.jl b/src/RingGrids/similar.jl deleted file mode 100644 index 5c95fed48..000000000 --- a/src/RingGrids/similar.jl +++ /dev/null @@ -1,24 +0,0 @@ -# todo this should be replaced by general broadcasting for <: AbstractGrid -for Grid in (:FullGaussianGrid, :FullClenshawGrid, :FullHEALPixGrid, :FullOctaHEALPixGrid, - :OctahedralGaussianGrid, :OctahedralClenshawGrid, :HEALPixGrid, - :OctaHEALPixGrid) - @eval begin - - function Base.similar(G::$Grid) - return $Grid{eltype(G)}(undef, G.nlat_half) - end - - function Base.similar(G::$Grid, ::Type{T}) where T - return $Grid{T}(undef, G.nlat_half) - end - - function Base.similar(G::$Grid, nlat_half::Integer) - return $Grid{eltype(G)}(undef, nlat_half) - end - - function Base.similar(G::$Grid, ::Type{T}, nlat_half::Integer) where T - return $Grid{T}(undef, nlat_half) - end - - end -end \ No newline at end of file diff --git a/src/RingGrids/utility_functions.jl b/src/RingGrids/utility_functions.jl index ab67da782..3d13fd3dd 100644 --- a/src/RingGrids/utility_functions.jl +++ b/src/RingGrids/utility_functions.jl @@ -1,8 +1,6 @@ -""" - true/false = isincreasing(v::Vector) - +"""$(TYPEDSIGNATURES) Check whether elements of a vector `v` are strictly increasing.""" -function isincreasing(x::Vector) +function isincreasing(x::AbstractVector) is_increasing = true for i in 2:length(x) is_increasing &= x[i-1] < x[i] ? true : false @@ -10,11 +8,9 @@ function isincreasing(x::Vector) return is_increasing end -""" - true/false = isdecreasing(v::Vector) - +"""$(TYPEDSIGNATURES) Check whether elements of a vector `v` are strictly decreasing.""" -function isdecreasing(x::Vector) +function isdecreasing(x::AbstractVector) is_decreasing = true for i in 2:length(x) is_decreasing &= x[i-1] > x[i] ? true : false @@ -22,22 +18,15 @@ function isdecreasing(x::Vector) return is_decreasing end -""" - true/false = extrema_in(v::Vector, a::Real, b::Real) - +"""$(TYPEDSIGNATURES) For every element vᵢ in v does a<=vi<=b hold?""" -function extrema_in(v::Vector, - a::Real, - b::Real) - +function extrema_in(v::AbstractVector, a::Real, b::Real) vmin, vmax = extrema(v) return (vmin >= a) && (vmax <= b) end # MATRIX rotations -""" - i_new, j_new = rotate_matrix_indices_90(i, j, s) - +"""$(TYPEDSIGNATURES) Rotate indices `i, j` of a square matrix of size s x s anti-clockwise by 90˚.""" @inline function rotate_matrix_indices_90(i::Integer, j::Integer, s::Integer) @boundscheck 0 < i <= s || throw(BoundsError) @@ -47,9 +36,7 @@ Rotate indices `i, j` of a square matrix of size s x s anti-clockwise by 90˚."" return i_new, j_new end -""" - i_new, j_new = rotate_matrix_indices_180(i, j, s) - +"""$(TYPEDSIGNATURES) Rotate indices `i, j` of a square matrix of size s x s by 180˚.""" @inline function rotate_matrix_indices_180(i::Integer, j::Integer, s::Integer) @boundscheck 0 < i <= s || throw(BoundsError) @@ -59,9 +46,7 @@ Rotate indices `i, j` of a square matrix of size s x s by 180˚.""" return i_new, j_new end -""" - i_new, j_new = rotate_matrix_indices_270(i, j, s) - +"""$(TYPEDSIGNATURES) Rotate indices `i, j` of a square matrix of size s x s anti-clockwise by 270˚.""" @inline function rotate_matrix_indices_270(i::Integer, j::Integer, s::Integer) @boundscheck 0 < i <= s || throw(BoundsError) diff --git a/src/SpeedyTransforms/spectral_transform.jl b/src/SpeedyTransforms/spectral_transform.jl index dcc36b731..0e9f4f50b 100644 --- a/src/SpeedyTransforms/spectral_transform.jl +++ b/src/SpeedyTransforms/spectral_transform.jl @@ -6,8 +6,8 @@ for the spectral transform.""" struct SpectralTransform{NF<:AbstractFloat} # GRID - Grid::Type{<:AbstractGrid} # grid type used - nlat_half::Int # resolution parameter of grid (# of latitudes on one hemisphere, Eq incl) + Grid::Type{<:AbstractGridArray} # grid type used + nlat_half::Int # resolution parameter of grid (# of latitudes on one hemisphere, Eq incl) # SPECTRAL RESOLUTION lmax::Int # Maximum degree l=[0, lmax] of spherical harmonics @@ -40,6 +40,7 @@ struct SpectralTransform{NF<:AbstractFloat} # SOLID ANGLES ΔΩ FOR QUADRATURE # (integration for the Legendre polynomials, extra normalisation of π/nlat included) + # vector is pole to pole although only northern hemisphere required solid_angles::Vector{NF} # = ΔΩ = sinθ Δθ Δϕ (solid angle of grid point) # RECURSION FACTORS @@ -70,26 +71,28 @@ Generator function for a SpectralTransform struct. With `NF` the number format, `Grid` the grid type `<:AbstractGrid` and spectral truncation `lmax, mmax` this function sets up necessary constants for the spetral transform. Also plans the Fourier transforms, retrieves the colatitudes, and preallocates the Legendre polynomials (if recompute_legendre == false) and quadrature weights.""" -function SpectralTransform( ::Type{NF}, # Number format NF - Grid::Type{<:AbstractGrid}, # type of spatial grid used - lmax::Int, # Spectral truncation: degrees - mmax::Int; # Spectral truncation: orders +function SpectralTransform( ::Type{NF}, # Number format NF + Grid::Type{<:AbstractGridArray}, # type of spatial grid used + lmax::Int, # Spectral truncation: degrees + mmax::Int; # Spectral truncation: orders recompute_legendre::Bool = true, # re or precompute legendre polynomials? legendre_shortcut::Symbol = :linear, # shorten Legendre loop over order m dealiasing::Real=DEFAULT_DEALIASING ) where NF + Grid = RingGrids.nonparametric_type(Grid) # always use nonparametric super type + # RESOLUTION PARAMETERS - nlat_half = get_nlat_half(mmax, dealiasing) # resolution parameter nlat_half, + nlat_half = get_nlat_half(mmax, dealiasing) # resolution parameter nlat_half, # number of latitude rings on one hemisphere incl equator - nlat = get_nlat(Grid, nlat_half) # 2nlat_half but one less if grids have odd # of lat rings - nlon_max = get_nlon_max(Grid, nlat_half) # number of longitudes around the equator + nlat = get_nlat(Grid, nlat_half) # 2nlat_half but one less if grids have odd # of lat rings + nlon_max = get_nlon_max(Grid, nlat_half) # number of longitudes around the equator # number of longitudes per latitude ring (one hemisphere only) nlons = [RingGrids.get_nlon_per_ring(Grid, nlat_half, j) for j in 1:nlat_half] nfreq_max = nlon_max÷2 + 1 # maximum number of fourier frequencies (real FFTs) # LATITUDE VECTORS (based on Gaussian, equi-angle or HEALPix latitudes) - colat = RingGrids.get_colat(Grid, nlat_half) # colatitude in radians + colat = RingGrids.get_colat(Grid, nlat_half) # colatitude in radians cos_colat = cos.(colat) # cos(colat) sin_colat = sin.(colat) # sin(colat) @@ -116,12 +119,12 @@ function SpectralTransform( ::Type{NF}, # Number format NF # LONGITUDE OFFSETS OF FIRST GRID POINT PER RING (0 for full and octahedral grids) _, lons = RingGrids.get_colatlons(Grid, nlat_half) - rings = eachring(Grid, nlat_half) # compute ring indices + rings = eachring(Grid, nlat_half) # compute ring indices lon1s = [lons[rings[j].start] for j in 1:nlat_half] # pick lons at first index for each ring lon_offsets = [cispi(m*lon1/π) for m in 0:mmax, lon1 in lon1s] # PREALLOCATE LEGENDRE POLYNOMIALS, +1 for 1-based indexing - Λ = zeros(LowerTriangularMatrix{NF}, lmax+1, mmax+1) # Legendre polynomials for one latitude + Λ = zeros(LowerTriangularMatrix{NF}, lmax+1, mmax+1) # Legendre polynomials for one latitude # allocate memory in Λs for polynomials at all latitudes or allocate dummy array if precomputed # Λs is of size (lmax+1) x (mmax+1) x nlat_half unless recomputed @@ -130,7 +133,7 @@ function SpectralTransform( ::Type{NF}, # Number format NF Λs = [zeros(LowerTriangularMatrix{NF}, b*(lmax+1), b*(mmax+1)) for _ in 1:nlat_half] if recompute_legendre == false # then precompute all polynomials - Λtemp = zeros(NF, lmax+1, mmax+1) # preallocate matrix + Λtemp = zeros(NF, lmax+1, mmax+1) # preallocate matrix for j in 1:nlat_half # only one hemisphere due to symmetry Legendre.λlm!(Λtemp, lmax, mmax, Float64(cos_colat[j])) # always precalculate in Float64 # underflow!(Λtemp, sqrt(floatmin(NF))) @@ -152,8 +155,8 @@ function SpectralTransform( ::Type{NF}, # Number format NF grad_x = [im*m for m in 0:mmax] # zonal gradient (precomputed currently not used) # meridional gradient for scalars (coslat scaling included) - grad_y1 = zeros(LowerTriangularMatrix, lmax+1, mmax+1) # term 1, mul with harmonic l-1, m - grad_y2 = zeros(LowerTriangularMatrix, lmax+1, mmax+1) # term 2, mul with harmonic l+1, m + grad_y1 = zeros(LowerTriangularMatrix, lmax+1, mmax+1) # term 1, mul with harmonic l-1, m + grad_y2 = zeros(LowerTriangularMatrix, lmax+1, mmax+1) # term 2, mul with harmonic l+1, m for m in 0:mmax # 0-based degree l, order m for l in m:lmax @@ -188,7 +191,7 @@ function SpectralTransform( ::Type{NF}, # Number format NF end end - vordiv_to_uv1[1, 1] = 0 # remove NaN from 0/0 + vordiv_to_uv1[1, 1] = 0 # remove NaN from 0/0 # EIGENVALUES (on unit sphere, hence 1/radius²-scaling is omitted) eigenvalues = [-l*(l+1) for l in 0:mmax+1] @@ -359,10 +362,10 @@ function gridded!( map::AbstractGrid{NF}, # gridded output acc_s = (acc_even - acc_odd) # and southern hemisphere # CORRECT FOR LONGITUDE OFFSETTS - o = lon_offsets[m, j_north] # longitude offset rotation + o = lon_offsets[m, j_north] # longitude offset rotation - gn[m] = muladd(acc_n, o, gn[m]) # accumulate in phase factors for northern - gs[m] = muladd(acc_s, o, gs[m]) # and southern hemisphere + gn[m] = muladd(acc_n, o, gn[m]) # accumulate in phase factors for northern + gs[m] = muladd(acc_s, o, gs[m]) # and southern hemisphere lm = lm_end + 1 # first index of next m column end @@ -419,8 +422,8 @@ function spectral!( alms::LowerTriangularMatrix{Complex{NF}}, # output: spectr @boundscheck get_nlat_half(map) == S.nlat_half || throw(BoundsError) # preallocate work warrays - fn = zeros(Complex{NF}, nfreq_max) # Fourier-transformed northern latitude - fs = zeros(Complex{NF}, nfreq_max) # Fourier-transformed southern latitude + fn = zeros(Complex{NF}, nfreq_max) # Fourier-transformed northern latitude + fs = zeros(Complex{NF}, nfreq_max) # Fourier-transformed southern latitude rings = eachring(map) # precompute ring indices @@ -460,7 +463,7 @@ function spectral!( alms::LowerTriangularMatrix{Complex{NF}}, # output: spectr an, as = fn[m], fs[m] # SOLID ANGLE QUADRATURE WEIGHTS and LONGITUDE OFFSET - o = lon_offsets[m, j_north] # longitude offset rotation + o = lon_offsets[m, j_north] # longitude offset rotation ΔΩ_rotated = ΔΩ*conj(o) # complex conjugate for rotation back to prime meridian # LEGENDRE TRANSFORM @@ -503,7 +506,7 @@ 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 + ) 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) @@ -540,12 +543,12 @@ end """ $(TYPEDSIGNATURES) Converts `map` to `Grid(map)` to execute `spectral(map::AbstractGrid; kwargs...)`.""" -function spectral( map::AbstractGrid{NF}; # gridded field - recompute_legendre::Bool = true, # saves memory - one_more_degree::Bool = false, # for lmax+2 x mmax+1 output size - ) where NF # number format NF +function spectral( map::AbstractGrid{NF}; # gridded field + recompute_legendre::Bool = true, # saves memory + one_more_degree::Bool = false, # for lmax+2 x mmax+1 output size + ) where NF # number format NF - Grid = RingGrids.nonparametric_type(typeof(map)) + Grid = typeof(map) trunc = get_truncation(map.nlat_half) S = SpectralTransform(NF, Grid, trunc+one_more_degree, trunc; recompute_legendre) return spectral(map, S) @@ -558,9 +561,9 @@ function spectral( map::AbstractGrid, # gridded field S::SpectralTransform{NF}, # spectral transform struct ) where NF # number format NF - map_NF = similar(map, NF) # convert map to NF + map_NF = similar(map, NF) # convert map to NF copyto!(map_NF, map) alms = LowerTriangularMatrix{Complex{NF}}(undef, S.lmax+1, S.mmax+1) - return spectral!(alms, map_NF, S) # in-place version + return spectral!(alms, map_NF, S) # in-place version end diff --git a/src/SpeedyWeather.jl b/src/SpeedyWeather.jl index bf7459578..2e4ea3d04 100644 --- a/src/SpeedyWeather.jl +++ b/src/SpeedyWeather.jl @@ -46,15 +46,15 @@ using .LowerTriangularMatrices # RingGrids export RingGrids -export FullClenshawGrid, - FullGaussianGrid, - FullHEALPixGrid, - FullOctaHEALPixGrid, - OctahedralGaussianGrid, - OctahedralClenshawGrid, - HEALPixGrid, - OctaHEALPixGrid, - plot +export FullClenshawGrid, FullClenshawArray, + FullGaussianGrid, FullGaussianArray, + FullHEALPixGrid, FullHEALPixArray, + FullOctaHEALPixGrid, FullOctaHEALPixArray, + OctahedralGaussianGrid, OctahedralGaussianArray, + OctahedralClenshawGrid, OctahedralClenshawArray, + HEALPixGrid, HEALPixArray, + OctaHEALPixGrid, OctaHEALPixArray, + eachring, eachgrid, plot include("RingGrids/RingGrids.jl") using .RingGrids diff --git a/src/dynamics/spectral_grid.jl b/src/dynamics/spectral_grid.jl index 646d66a4e..af6c15c1c 100644 --- a/src/dynamics/spectral_grid.jl +++ b/src/dynamics/spectral_grid.jl @@ -75,7 +75,7 @@ function Base.show(io::IO, SG::SpectralGrid) (; NF, trunc, Grid, radius, nlat, npoints, nlev, vertical_coordinates) = SG (; n_particles) = SG - # resolution information + # resolution information res_ave = sqrt(4π*radius^2/npoints)/1000 # in [km] s(x) = x > 1000 ? @sprintf("%i", x) : @sprintf("%.3g", x) diff --git a/src/output/output.jl b/src/output/output.jl index c946c2d51..0681fda35 100644 --- a/src/output/output.jl +++ b/src/output/output.jl @@ -38,7 +38,7 @@ Base.@kwdef mutable struct OutputWriter{NF<:Union{Float32, Float64}, Model<:Mode spectral_grid::SpectralGrid - # FILE OPTIONS + # FILE OPTIONS output::Bool = false "[OPTION] path to output folder, run_???? will be created within" @@ -88,7 +88,7 @@ Base.@kwdef mutable struct OutputWriter{NF<:Union{Float32, Float64}, Model<:Mode # INPUT GRID (the one used in the dynamical core) input_Grid::Type{<:AbstractGrid} = spectral_grid.Grid - # Output as matrix (particularly for reduced grids) + # Output as matrix (particularly for reduced grids) "[OPTION] sort grid points into a matrix (interpolation-free), for OctahedralClenshawGrid, OctaHEALPixGrid only" as_matrix::Bool = false @@ -98,7 +98,7 @@ Base.@kwdef mutable struct OutputWriter{NF<:Union{Float32, Float64}, Model<:Mode # OUTPUT GRID "[OPTION] the grid used for output, full grids only" - output_Grid::Type{<:AbstractFullGrid} = RingGrids.full_grid(input_Grid) + output_Grid::Type{<:AbstractFullGrid} = RingGrids.full_grid_type(input_Grid) "[OPTION] the resolution of the output grid, default: same nlat_half as in the dynamical core" nlat_half::Int = spectral_grid.nlat_half @@ -175,7 +175,7 @@ function initialize!( feedback.progress_meter.desc = "Weather is speedy: run $(output.id) " feedback.output = true # if output=true set feedback.output=true too! - # OUTPUT FREQUENCY + # OUTPUT FREQUENCY output.output_every_n_steps = max(1, round(Int, Millisecond(output.output_dt).value/time_stepping.Δt_millisec.value)) output.output_dt = Second(round(Int, output.output_every_n_steps*time_stepping.Δt_sec)) diff --git a/test/grids.jl b/test/grids.jl index 51a0f15f4..69b7cad2a 100644 --- a/test/grids.jl +++ b/test/grids.jl @@ -1,21 +1,66 @@ +using JLArrays +using Adapt + +@testset "Grid types" begin + for G in ( + FullClenshawGrid, + FullGaussianGrid, + OctahedralGaussianGrid, + OctahedralClenshawGrid, + HEALPixGrid, + OctaHEALPixGrid, + FullHEALPixGrid, + FullOctaHEALPixGrid + ) + + full = RingGrids.full_grid_type(G) + @test full == RingGrids.full_grid_type(G) + @test full <: RingGrids.AbstractFullGrid + if ~(G <: RingGrids.AbstractFullGrid) + @test G <: RingGrids.AbstractReducedGrid + end + end + + for G in ( + FullClenshawArray, + FullGaussianArray, + OctahedralGaussianArray, + OctahedralClenshawArray, + HEALPixArray, + OctaHEALPixArray, + FullHEALPixArray, + FullOctaHEALPixArray + ) + + full = RingGrids.full_grid_type(G) + @test full == RingGrids.full_grid_type(G) + @test full == RingGrids.horizontal_grid_type(RingGrids.full_array_type(G)) + @test RingGrids.full_array_type(G) == RingGrids.full_array_type(RingGrids.horizontal_grid_type(G)) + @test full <: RingGrids.AbstractFullGridArray + if ~(G <: RingGrids.AbstractFullGridArray) + @test G <: RingGrids.AbstractReducedGridArray + end + end +end + @testset "Grid indexing" begin for NF in (Float32, Float64) # with vector and resolution parameter provided - L = FullClenshawGrid(randn(NF, 96*47), 24) # L24 grid - F = FullGaussianGrid(randn(NF, 96*48), 24) # F24 grid - O = OctahedralGaussianGrid(randn(NF, 3168), 24) # O24 grid - C = OctahedralClenshawGrid(randn(NF, 3056), 24) # C24 grid - H = HEALPixGrid(randn(NF, 3072), 32) # H32 grid - J = OctaHEALPixGrid(randn(NF, 4096), 32) # J32 grid - K = FullOctaHEALPixGrid(randn(NF, 128*63), 32) # K32 grid + L = FullClenshawGrid(randn(NF, 96*47), 24) # L24 grid + F = FullGaussianGrid(randn(NF, 96*48), 24) # F24 grid + O = OctahedralGaussianGrid(randn(NF, 3168), 24) # O24 grid + C = OctahedralClenshawGrid(randn(NF, 3056), 24) # C24 grid + H = HEALPixGrid(randn(NF, 3072), 32) # H32 grid + J = OctaHEALPixGrid(randn(NF, 4096), 32) # J32 grid + K = FullOctaHEALPixGrid(randn(NF, 128*63), 32) # K32 grid # without resolution parameter provided (inferred from vector length) - L2 = FullClenshawGrid(randn(NF, 96*47)) # L24 grid - F2 = FullGaussianGrid(randn(NF, 96*48)) # F24 grid - O2 = OctahedralGaussianGrid(randn(NF, 3168)) # O24 grid - C2 = OctahedralClenshawGrid(randn(NF, 3056)) # C24 grid - H2 = HEALPixGrid(randn(NF, 3072)) # H32 grid + L2 = FullClenshawGrid(randn(NF, 96*47)) # L24 grid + F2 = FullGaussianGrid(randn(NF, 96*48)) # F24 grid + O2 = OctahedralGaussianGrid(randn(NF, 3168)) # O24 grid + C2 = OctahedralClenshawGrid(randn(NF, 3056)) # C24 grid + H2 = HEALPixGrid(randn(NF, 3072)) # H32 grid J2 = OctaHEALPixGrid(randn(NF, 4096)) # J32 grid K2 = FullOctaHEALPixGrid(randn(NF, 128*63)) # K32 grid @@ -204,8 +249,8 @@ end grid = zeros(G, n) # precompute indices and boundscheck - rings = SpeedyWeather.eachring(grid) - rings2 = [SpeedyWeather.each_index_in_ring(grid, j) for j in 1:SpeedyWeather.get_nlat(grid)] + rings = RingGrids.eachring(grid) + rings2 = [RingGrids.each_index_in_ring(grid, j) for j in 1:RingGrids.get_nlat(grid)] @test rings == rings2 end @@ -214,14 +259,14 @@ end @testset "Grid broadcasting" begin n = 2 - @testset for G in ( FullClenshawGrid, - FullGaussianGrid, - OctahedralGaussianGrid, - OctahedralClenshawGrid, - HEALPixGrid, - OctaHEALPixGrid, - FullHEALPixGrid, - FullOctaHEALPixGrid, + @testset for G in ( FullClenshawArray, + FullGaussianArray, + OctahedralGaussianArray, + OctahedralClenshawArray, + HEALPixArray, + OctaHEALPixArray, + FullHEALPixArray, + FullOctaHEALPixArray, ) @test zeros(G, n) .+ 1 == ones(G, n) @@ -230,6 +275,18 @@ end @test zeros(G, n) + ones(G, n) == ones(G, n) @test 2ones(G, n) == ones(G, n) + ones(G, n) + # don't promote to Array + for s in ((n,), (n, n), (n, n, n), (n, n, n, n)) + grid = zeros(G, s...) + @test (grid + grid) isa G + @test (grid .+ grid) isa G + @test (grid - grid) isa G + @test (grid .- grid) isa G + @test (grid .* grid) isa G + @test (grid ./ grid) isa G + @test 2grid isa G + end + # promote types, Grid{Float16} -> Grid{Float64} etc @test all(ones(G{Float16}, n)*2.0 .=== 2.0) @test all(ones(G{Float16}, n)*2f0 .=== 2f0) @@ -255,4 +312,123 @@ end @test all(ones(G{Float16}, n) ./ ones(G{Float64}, n) .=== 1.0) @test all(ones(G{Float32}, n) ./ ones(G{Float64}, n) .=== 1.0) end +end + +@testset "N-dimensional indexing" begin + m, n, p = 2, 3, 4 + @testset for G in ( FullClenshawGrid, + FullGaussianGrid, + OctahedralGaussianGrid, + OctahedralClenshawGrid, + HEALPixGrid, + OctaHEALPixGrid, + FullHEALPixGrid, + FullOctaHEALPixGrid, + ) + + grid = rand(G, m, n, p) + @test grid[:, 1, 1] isa G + @test grid[1] == grid.data[1] + @test grid[1, 1, 1] == grid.data[1, 1, 1] + + @test grid[1:2, 1:2, 1:2] == grid.data[1:2, 1:2, 1:2] + @test grid[1, 1, :] == grid.data[1, 1, :] + + idx = CartesianIndex((1, 2, 3)) + @test grid[idx] == grid.data[idx] + + ids = CartesianIndices((m, n, p)) + @test grid[ids] == grid.data[ids] + @test grid[ids] isa Array + end +end + +@testset "Loop indexing" begin + n = 2 + @testset for G in ( FullClenshawArray, + FullGaussianArray, + OctahedralGaussianArray, + OctahedralClenshawArray, + HEALPixArray, + OctaHEALPixArray, + FullHEALPixArray, + FullOctaHEALPixArray, + ) + + for s in ((n,), (n, n), (n, n, n), (n, n, n, n)) + grid = zeros(G, s...) + + for k in eachgrid(grid) + for (j, ring) in enumerate(eachring(grid)) + for ij in ring + grid[ij, k] = 1 + end + end + end + @test all(grid .== 1) + end + end +end + +@testset "AbstractGridArray: GPU (JLArrays)" begin + NF = Float32 + @testset for Grid in ( + FullClenshawArray, + FullGaussianArray, + OctahedralGaussianArray, + OctahedralClenshawArray, + HEALPixArray, + OctaHEALPixArray, + FullHEALPixArray, + FullOctaHEALPixArray, + ) + s = (2, 3, 4) + ndims = length(s) + + G_cpu = randn(Grid{NF}, s...) + + # constructors/adapt + G = adapt(JLArray, G_cpu) + G2 = Grid(adapt(JLArray, G_cpu.data)) + @test G == G2 + + # broadcasting doesn't escape + @test G + G isa Grid{NF, ndims, JLArray{NF, ndims}} + @test G .+ G isa Grid{NF, ndims, JLArray{NF, ndims}} + @test G_cpu + G_cpu isa Grid{NF, ndims, Array{NF, ndims}} + @test G_cpu .+ G_cpu isa Grid{NF, ndims, Array{NF, ndims}} + + # getindex + @test G[1, :, :] isa JLArray{NF, 2} + @test G[:, 1, 1] isa Grid{NF, 1, JLArray{NF, 1}} + for k in eachgrid(G) + for (j, ring) in enumerate(eachring(G)) + @test G[ring, k] == adapt(JLArray, G_cpu[ring, k]) + @test Array(G[ring, k]) == G_cpu[ring, k] + end + end + + # setindex! + G_test = JLArray(rand(NF,s[3])) + G[1, 1, :] .= G_test # with . + @test G[1, 1, :] == G_test + + G_test = JLArray(rand(NF,s[3])) + G[1, 1, :] = G_test # without . + @test G[1, 1, :] == G_test + + # with other grid {Array} + G_test = rand(Grid, s[1]) + G[:, 1, 1] = G_test # conversion to float64 -> float32 + @test Array(G[:, 1, 1].data) ≈ G_test + + # with other grid {JLArray} + G_test = adapt(JLArray,rand(Grid, s[1])) + G[:, 1, 1] = G_test + @test G[:, 1, 1].data ≈ G_test.data + + # fill + fill!(G, 2) + @test all(G .== 2) + end end \ No newline at end of file diff --git a/test/lower_triangular_matrix.jl b/test/lower_triangular_matrix.jl index ae11ac81a..3c0a4075c 100644 --- a/test/lower_triangular_matrix.jl +++ b/test/lower_triangular_matrix.jl @@ -1,8 +1,4 @@ -using JLArrays - -# remove with extension -LowerTriangularMatrices.nonparametric_type(::Type{<:JLArray}) = JLArray -using Adapt +using JLArrays, Adapt import Random @testset "LowerTriangularMatrix" begin