From 2e436d5c295b7d4673d0b414314b0e8efc17e43e Mon Sep 17 00:00:00 2001 From: Milan Date: Wed, 10 Apr 2024 17:13:05 -0400 Subject: [PATCH 01/35] First ideas for AbstractGridArray --- src/RingGrids/RingGrids.jl | 13 +++--- src/RingGrids/arrays/full_gaussian.jl | 25 ++++++++++++ .../{grids_general.jl => general.jl} | 40 ++++++++++++------- src/RingGrids/{ => grids}/full_grids.jl | 0 src/RingGrids/{ => grids}/healpix.jl | 0 src/RingGrids/{ => grids}/octahealpix.jl | 0 src/RingGrids/{ => grids}/octahedral.jl | 0 7 files changed, 56 insertions(+), 22 deletions(-) create mode 100644 src/RingGrids/arrays/full_gaussian.jl rename src/RingGrids/{grids_general.jl => general.jl} (88%) rename src/RingGrids/{ => grids}/full_grids.jl (100%) rename src/RingGrids/{ => grids}/healpix.jl (100%) rename src/RingGrids/{ => grids}/octahealpix.jl (100%) rename src/RingGrids/{ => grids}/octahedral.jl (100%) diff --git a/src/RingGrids/RingGrids.jl b/src/RingGrids/RingGrids.jl index 4dc5b857b..fb2496984 100644 --- a/src/RingGrids/RingGrids.jl +++ b/src/RingGrids/RingGrids.jl @@ -67,15 +67,12 @@ export interpolate, update_locator, update_locator! -# export plot - include("utility_functions.jl") - -include("grids_general.jl") -include("full_grids.jl") -include("octahedral.jl") -include("healpix.jl") -include("octahealpix.jl") +include("general.jl") +include("grids/full_grids.jl") +include("grids/octahedral.jl") +include("grids/healpix.jl") +include("grids/octahealpix.jl") include("quadrature_weights.jl") include("interpolation.jl") include("show.jl") diff --git a/src/RingGrids/arrays/full_gaussian.jl b/src/RingGrids/arrays/full_gaussian.jl new file mode 100644 index 000000000..6ee946bc7 --- /dev/null +++ b/src/RingGrids/arrays/full_gaussian.jl @@ -0,0 +1,25 @@ +struct FullGaussianArray{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? +end + +const FullGaussianGrid{T} = FullGaussianArray{T, 1, Vector{T}} + + + +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 \ No newline at end of file diff --git a/src/RingGrids/grids_general.jl b/src/RingGrids/general.jl similarity index 88% rename from src/RingGrids/grids_general.jl rename to src/RingGrids/general.jl index b7dcae803..521f19611 100644 --- a/src/RingGrids/grids_general.jl +++ b/src/RingGrids/general.jl @@ -1,29 +1,41 @@ -abstract type AbstractGrid{T} <: AbstractVector{T} end +abstract type AbstractGridArray{T, N, ArrayType <: AbstractArray{T, N}} <: AbstractArray{T,N} end +abstract type AbstractFullGridArray{T, N, ArrayType <: AbstractArray{T, N}} <: AbstractGridArray{T,N} end +abstract type AbstractReducedGridArray{T, N, ArrayType <: AbstractArray{T, N}} <: AbstractGridArray{T,N} 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 +Base.length(G::AbstractGridArray) = length(G.data) +Base.size(G::AbstractGridArray) = size(G.data) +Base.sizeof(G::AbstractGridArray) = sizeof(G.data) + +@inline Base.getindex(G::AbstractGridArray, ijk::Integer...) = getindex(G.data,ijk...) +@inline Base.getindex(G::AbstractGridArray, r::AbstractRange) = get_index(G.data, r) + + + +@inline function Base.getindex(L::AbstractGridArray, col::Colon, I...) + return AbstractGridArray(getindex(L.data, col, I...), L.m, L.n) 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.getindex(L::AbstractGridArray, r::AbstractRange) = getindex(L.data,r) +@inline Base.getindex(L::AbstractGridArray, r::AbstractRange, I...) = getindex(L.data, r, I...) +@inline Base.getindex(L::AbstractGridArray, i::Integer) = getindex(L.data,i) +@inline Base.getindex(L::AbstractGridArray, I::CartesianIndex) = getindex(L, Tuple(I)...) + + @inline Base.setindex!(G::AbstractGrid, x::AbstractVector, r::AbstractRange) = setindex!(G.data, x, r) +@inline Base.setindex!(G::AbstractGrid, x, k::Integer) = setindex!(G.data, x, k) + -function grids_match(A::AbstractGrid, B::AbstractGrid) +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{<:AbstractGrid}, B::Type{<:AbstractGrid}) - # throws an error for non-parametric types... - return A.name.wrapper == B.name.wrapper +function _grids_match(A::Type{<:AbstractGridArray}, B::Type{<:AbstractGridArray}) + return nonparametric_type(A) == nonparametric_type(B) end """ diff --git a/src/RingGrids/full_grids.jl b/src/RingGrids/grids/full_grids.jl similarity index 100% rename from src/RingGrids/full_grids.jl rename to src/RingGrids/grids/full_grids.jl diff --git a/src/RingGrids/healpix.jl b/src/RingGrids/grids/healpix.jl similarity index 100% rename from src/RingGrids/healpix.jl rename to src/RingGrids/grids/healpix.jl diff --git a/src/RingGrids/octahealpix.jl b/src/RingGrids/grids/octahealpix.jl similarity index 100% rename from src/RingGrids/octahealpix.jl rename to src/RingGrids/grids/octahealpix.jl diff --git a/src/RingGrids/octahedral.jl b/src/RingGrids/grids/octahedral.jl similarity index 100% rename from src/RingGrids/octahedral.jl rename to src/RingGrids/grids/octahedral.jl From 2c1786ebae12575370567cab11d4f7413fbd3954 Mon Sep 17 00:00:00 2001 From: Milan Date: Fri, 12 Apr 2024 19:24:13 -0400 Subject: [PATCH 02/35] AbstractGridArray constructors --- src/RingGrids/RingGrids.jl | 32 +- src/RingGrids/arrays/full_gaussian.jl | 25 -- src/RingGrids/full_grids.jl | 86 +++++ src/RingGrids/general.jl | 353 +++++++++--------- src/RingGrids/grids/full_clenshaw.jl | 31 ++ src/RingGrids/grids/full_gaussian.jl | 35 ++ src/RingGrids/grids/full_grids.jl | 204 ---------- src/RingGrids/grids/full_healpix.jl | 30 ++ src/RingGrids/grids/full_octahealpix.jl | 30 ++ src/RingGrids/grids/new_grids.jl | 42 +++ .../{octahedral.jl => octahedral_clenshaw.jl} | 0 src/RingGrids/grids/octahedral_gaussian.jl | 0 src/RingGrids/interpolation.jl | 12 +- src/RingGrids/reduced_grids.jl | 4 + src/RingGrids/scaling.jl | 18 +- src/RingGrids/utility_functions.jl | 33 +- 16 files changed, 489 insertions(+), 446 deletions(-) delete mode 100644 src/RingGrids/arrays/full_gaussian.jl create mode 100644 src/RingGrids/full_grids.jl create mode 100644 src/RingGrids/grids/full_clenshaw.jl create mode 100644 src/RingGrids/grids/full_gaussian.jl delete mode 100644 src/RingGrids/grids/full_grids.jl create mode 100644 src/RingGrids/grids/full_healpix.jl create mode 100644 src/RingGrids/grids/full_octahealpix.jl create mode 100644 src/RingGrids/grids/new_grids.jl rename src/RingGrids/grids/{octahedral.jl => octahedral_clenshaw.jl} (100%) create mode 100644 src/RingGrids/grids/octahedral_gaussian.jl create mode 100644 src/RingGrids/reduced_grids.jl diff --git a/src/RingGrids/RingGrids.jl b/src/RingGrids/RingGrids.jl index fb2496984..967073e32 100644 --- a/src/RingGrids/RingGrids.jl +++ b/src/RingGrids/RingGrids.jl @@ -68,15 +68,31 @@ export interpolate, update_locator! include("utility_functions.jl") + +# GENERAL include("general.jl") -include("grids/full_grids.jl") -include("grids/octahedral.jl") -include("grids/healpix.jl") -include("grids/octahealpix.jl") -include("quadrature_weights.jl") -include("interpolation.jl") -include("show.jl") -include("similar.jl") +include("full_grids.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") +# include("similar.jl") + +# OUTPUT +include("show.jl") + end \ No newline at end of file diff --git a/src/RingGrids/arrays/full_gaussian.jl b/src/RingGrids/arrays/full_gaussian.jl deleted file mode 100644 index 6ee946bc7..000000000 --- a/src/RingGrids/arrays/full_gaussian.jl +++ /dev/null @@ -1,25 +0,0 @@ -struct FullGaussianArray{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? -end - -const FullGaussianGrid{T} = FullGaussianArray{T, 1, Vector{T}} - - - -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 \ No newline at end of file diff --git a/src/RingGrids/full_grids.jl b/src/RingGrids/full_grids.jl new file mode 100644 index 000000000..561dfcf4c --- /dev/null +++ b/src/RingGrids/full_grids.jl @@ -0,0 +1,86 @@ +abstract type AbstractFullGridArray{T, N, ArrayType <: AbstractArray{T, N}} <: AbstractGridArray{T, N, ArrayType} end + +""" +abstract type AbstractFullGrid{T} <: AbstractGrid{T} end + +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.""" +const AbstractFullGrid{T} = AbstractFullGridArray{T, 1, Vector{T}} + +full_grid(G::Type{<:AbstractFullGridArray}) = G + + +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(Grid, nlat_half) + +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)) +end + +function get_lond(Grid::Type{<:AbstractFullGrid}, 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) + + 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 + colats = zeros(npoints) # preallocate + lons = zeros(npoints) + + for j in 1:nlat # populate preallocated colats, lons + for i in 1:nlon + ij = i + (j-1)*nlon # continuous index ij + colats[ij] = colat[j] + lons[ij] = lon[i] + end + end + + return colats, lons +end + +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 = 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{<:AbstractFullGridArray}, + 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 \ No newline at end of file diff --git a/src/RingGrids/general.jl b/src/RingGrids/general.jl index 521f19611..7db89c2a7 100644 --- a/src/RingGrids/general.jl +++ b/src/RingGrids/general.jl @@ -1,195 +1,183 @@ -abstract type AbstractGridArray{T, N, ArrayType <: AbstractArray{T, N}} <: AbstractArray{T,N} end -abstract type AbstractFullGridArray{T, N, ArrayType <: AbstractArray{T, N}} <: AbstractGridArray{T,N} end -abstract type AbstractReducedGridArray{T, N, ArrayType <: AbstractArray{T, N}} <: AbstractGridArray{T,N} end +abstract type AbstractGridArray{T, N, ArrayType <: AbstractArray{T, N}} <: AbstractArray{T, N} end +const AbstractGrid{T} = AbstractGridArray{T, 1, Vector{T}} -# all AbstractGrids have their grid points stored in a vector field `data` -# propagate length, size, getindex, setindex! for that +# nonparametric_type(G::Type{<:AbstractGridArray}) = @warn "Please define nonparametric_type(::$G)" +# horizontal_grid_type(G::Type{<:AbstractGridArray}) = @warn "Please define horizontal_grid_type(::$G)" + +# LENGTH SIZE Base.length(G::AbstractGridArray) = length(G.data) Base.size(G::AbstractGridArray) = size(G.data) Base.sizeof(G::AbstractGridArray) = sizeof(G.data) -@inline Base.getindex(G::AbstractGridArray, ijk::Integer...) = getindex(G.data,ijk...) -@inline Base.getindex(G::AbstractGridArray, r::AbstractRange) = get_index(G.data, r) - - - -@inline function Base.getindex(L::AbstractGridArray, col::Colon, I...) - return AbstractGridArray(getindex(L.data, col, I...), L.m, L.n) -end - - +get_nlat_half(grid::AbstractGridArray) = grid.nlat_half +nlat_odd(grid::AbstractGridArray) = nlat_odd(typeof(grid)) -@inline Base.getindex(L::AbstractGridArray, r::AbstractRange) = getindex(L.data,r) -@inline Base.getindex(L::AbstractGridArray, r::AbstractRange, I...) = getindex(L.data, r, I...) -@inline Base.getindex(L::AbstractGridArray, i::Integer) = getindex(L.data,i) -@inline Base.getindex(L::AbstractGridArray, I::CartesianIndex) = getindex(L, Tuple(I)...) - - -@inline Base.setindex!(G::AbstractGrid, x::AbstractVector, r::AbstractRange) = setindex!(G.data, x, r) -@inline Base.setindex!(G::AbstractGrid, x, k::Integer) = setindex!(G.data, x, k) +# 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) +get_nlat(grid::Grid) where {Grid<:AbstractGridArray} = get_nlat(Grid, grid.nlat_half) +# get total number of grid points vs horizontal (2D) only +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) +get_npoints2D(grid::Grid) where {Grid<:AbstractGridArray} = get_npoints2D(Grid, grid.nlat_half) -function grids_match(A::AbstractGridArray, B::AbstractGridArray) - length(A) == length(B) && return _grids_match(typeof(A), typeof(B)) - return false +## INDEXING +@inline Base.getindex(G::AbstractGridArray, ijk::Integer...) = getindex(G.data,ijk...) +@inline Base.getindex(G::AbstractGridArray, r::AbstractRange, k::Integer...) = getindex(G.data, r, k...) +@inline Base.getindex(G::AbstractGridArray, ij::Integer, k::AbstractRange) = getindex(G.data, ij, k) +@inline Base.getindex(G::AbstractGridArray, I::CartesianIndex) = getindex(G, Tuple(I)...) + +@inline function Base.getindex( + G::GridArray, + col::Colon, + k..., +) where {GridArray<:AbstractGridArray{T, N, ArrayType}} where {T, N, ArrayType} + GridArray_ = nonparametric_type(GridArray) + return GridArray_{T, 1}(getindex(G.data, col, k...), G.nlat_half) end -function _grids_match(A::Type{<:AbstractGridArray}, B::Type{<:AbstractGridArray}) - return nonparametric_type(A) == nonparametric_type(B) +@inline Base.setindex!(G::AbstractGridArray, x, ijk::Integer...) = + setindex!(G.data, x, ijk...) +@inline Base.setindex!(G::AbstractGridArray, x::AbstractVector, ij::AbstractRange, k::Integer...) = + setindex!(G.data, x, ij, k...) +@inline Base.setindex!(G::AbstractGridArray, x::AbstractVector, ij::Integer, k::AbstractRange) = + setindex!(G.data, x, ij, k) + +## CONSTRUCTORS +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 -""" - 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) +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 - -`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) +# if no rings are provided, calculate them +function (::Type{Grid})(data::AbstractArray, nlat_half::Integer) where {Grid<:AbstractGridArray} + GridArray_ = nonparametric_type(Grid) + rings = eachring(Grid, nlat_half) + return GridArray_(data, nlat_half, rings) 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) +# if no nlat_half provided calculate it +function (::Type{Grid})(data::AbstractArray) where {Grid<:AbstractGridArray} + npoints2D = size(data, 1) + nlat_half = get_nlat_half(Grid, npoints2D) + return Grid(data, nlat_half) +end -# 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) +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 -# 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 +# zero element of an AbstractGridArray instance grid by creating new zero(grid.data) +Base.zero(grid::Grid) where {Grid<:AbstractGridArray} = Grid(zero(grid.data), grid.nlat_half, grid.rings) -# does the grid have an odd number of latitudes? -nlat_odd(grid::AbstractGrid) = nlat_odd(typeof(grid)) +# 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} + return Grid(ArrayType{T, N}(undef, get_npoints2D(Grid, nlat_half), k...), nlat_half) +end -# 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) +# 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 -# get total number of grid pointst -get_npoints(grid::Grid) where {Grid<:AbstractGrid} = get_npoints(Grid, grid.nlat_half) +# 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 -get_latdlonds(grid::Grid) where {Grid<:AbstractGrid} = get_latdlonds(Grid, grid.nlat_half) +## COORDINATES +get_latdlonds(grid::Grid) where {Grid<:AbstractGridArray} = 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 +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 .= latds .* (360/2π) # now in degrees 90˚...-90˚ + 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 +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 -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) +get_lat(grid::Grid) where {Grid<:AbstractGridArray} = get_lat(Grid, grid.nlat_half) +get_latd(grid::Grid) where {Grid<:AbstractGridArray} = get_latd(Grid, grid.nlat_half) +get_lond(grid::Grid) where {Grid<:AbstractGridArray} = get_lond(Grid, grid.nlat_half) +get_lon(grid::Grid) where {Grid<:AbstractGridArray} = get_lon(Grid, grid.nlat_half) +get_colat(grid::Grid) where {Grid<:AbstractGridArray} = get_colat(Grid, grid.nlat_half) +get_colatlons(grid::Grid) where {Grid<:AbstractGridArray} = get_colatlons(Grid, grid.nlat_half) -function get_lat(Grid::Type{<:AbstractGrid}, nlat_half::Integer) +function get_lat(Grid::Type{<:AbstractGridArray}, nlat_half::Integer) return π/2 .- get_colat(Grid, nlat_half) end -function get_latd(Grid::Type{<:AbstractGrid}, nlat_half::Integer) +function get_latd(Grid::Type{<:AbstractGridArray}, 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 +get_lon(::Type{<:AbstractGridArray}, nlat_half::Integer) = Float64[] +get_lond(::Type{<:AbstractGridArray}, nlat_half::Integer) = Float64[] +## ITERATORS """ $(TYPEDSIGNATURES) Vector{UnitRange} `rings` to loop over every ring of grid `grid` @@ -199,9 +187,9 @@ and then each grid point per ring. To be used like for ring in rings for ij in ring grid[ij]""" -eachring(grid::AbstractGrid) = eachring(typeof(grid), grid.nlat_half) +eachring(grid::AbstractGridArray) = grid.rings -function eachring(Grid::Type{<:AbstractGrid}, nlat_half::Integer) +function eachring(Grid::Type{<:AbstractGridArray}, nlat_half::Integer) rings = Vector{UnitRange{Int}}(undef, get_nlat(Grid, nlat_half)) each_index_in_ring!(rings, Grid, nlat_half) return rings @@ -211,13 +199,39 @@ 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} +function eachring(grid1::Grid, grids::Grid...) where {Grid<:AbstractGridArray} @inline n = length(grid1) Base._all_match_first(X->length(X), n, grid1, grids...) || throw(BoundsError) return eachring(grid1) end +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}) + 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 Vector{UnitRange} describing rind indices @@ -236,31 +250,34 @@ $(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) +function get_nlons(Grid::Type{<:AbstractGridArray}, 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) +get_nlon_max(grid::Grid) where {Grid<:AbstractGridArray} = 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 +# # following https://docs.julialang.org/en/v1/manual/interfaces/#man-interfaces-broadcasting +# import Base.Broadcast: BroadcastStyle, Broadcasted -# 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)}() +# # {1} as grids are <:AbstractVector, Grid here is the non-parameteric Grid type! +# struct AbstractGridArrayStyle{Grid} <: Broadcast.AbstractArrayStyle{1} end -# 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 +# # 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} = +# AbstractGridArrayStyle{nonparametric_type(Grid)}() + +# # allocation for broadcasting, create a new Grid with undef of type/number format T +# function Base.similar(bc::Broadcasted{AbstractGridArrayStyle{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 +# # ::Val{0} for broadcasting with 0-dimensional, ::Val{1} for broadcasting with vectors, etc +# AbstractGridArrayStyle{Grid}(::Val{0}) where Grid = AbstractGridArrayStyle{Grid}() +# AbstractGridArrayStyle{Grid}(::Val{1}) where Grid = AbstractGridArrayStyle{Grid}() +# AbstractGridArrayStyle{Grid}(::Val{2}) where Grid = AbstractGridArrayStyle{Grid}() +# AbstractGridArrayStyle{Grid}(::Val{3}) where Grid = AbstractGridArrayStyle{Grid}() \ 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..68c96db49 --- /dev/null +++ b/src/RingGrids/grids/full_clenshaw.jl @@ -0,0 +1,31 @@ +""" + 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 \ 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..02448c5c0 --- /dev/null +++ b/src/RingGrids/grids/full_gaussian.jl @@ -0,0 +1,35 @@ +struct FullGaussianArray{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? + + 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 +horizontal_grid_type(::Type{<:FullGaussianArray}) = FullGaussianGrid + +# SIZE +nlat_odd(::Type{<:FullGaussianArray}) = false +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_grids.jl b/src/RingGrids/grids/full_grids.jl deleted file mode 100644 index 84ab8d904..000000000 --- a/src/RingGrids/grids/full_grids.jl +++ /dev/null @@ -1,204 +0,0 @@ -""" - abstract type AbstractFullGrid{T} <: AbstractGrid{T} end - -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 - -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) - -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)) -end - -function get_lond(Grid::Type{<:AbstractFullGrid}, 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) - - 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 - colats = zeros(npoints) # preallocate - lons = zeros(npoints) - - for j in 1:nlat # populate preallocated colats, lons - for i in 1:nlon - ij = i + (j-1)*nlon # continuous index ij - colats[ij] = colat[j] - lons[ij] = lon[i] - end - 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 diff --git a/src/RingGrids/grids/full_healpix.jl b/src/RingGrids/grids/full_healpix.jl new file mode 100644 index 000000000..2ea857c47 --- /dev/null +++ b/src/RingGrids/grids/full_healpix.jl @@ -0,0 +1,30 @@ +""" + 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) + diff --git a/src/RingGrids/grids/full_octahealpix.jl b/src/RingGrids/grids/full_octahealpix.jl new file mode 100644 index 000000000..76348dee4 --- /dev/null +++ b/src/RingGrids/grids/full_octahealpix.jl @@ -0,0 +1,30 @@ +""" + 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 diff --git a/src/RingGrids/grids/new_grids.jl b/src/RingGrids/grids/new_grids.jl new file mode 100644 index 000000000..f2cfc36dc --- /dev/null +++ b/src/RingGrids/grids/new_grids.jl @@ -0,0 +1,42 @@ +""" + 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 +""" \ No newline at end of file diff --git a/src/RingGrids/grids/octahedral.jl b/src/RingGrids/grids/octahedral_clenshaw.jl similarity index 100% rename from src/RingGrids/grids/octahedral.jl rename to src/RingGrids/grids/octahedral_clenshaw.jl diff --git a/src/RingGrids/grids/octahedral_gaussian.jl b/src/RingGrids/grids/octahedral_gaussian.jl new file mode 100644 index 000000000..e69de29bb diff --git a/src/RingGrids/interpolation.jl b/src/RingGrids/interpolation.jl index 3ca9fa0fb..bacbcdc60 100644 --- a/src/RingGrids/interpolation.jl +++ b/src/RingGrids/interpolation.jl @@ -213,9 +213,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 +315,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 +511,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 +522,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/reduced_grids.jl b/src/RingGrids/reduced_grids.jl new file mode 100644 index 000000000..97b0da886 --- /dev/null +++ b/src/RingGrids/reduced_grids.jl @@ -0,0 +1,4 @@ +abstract type AbstractReducedGridArray{T, N, ArrayType <: AbstractArray{T, N}} <: AbstractGridArray{T, N, ArrayType} end +const AbstractReducedGrid{T} = AbstractReducedGridArray{T, 1, Vector{T}} + +full_grid(G::Type{<:AbstractReducedGridArray}) = @warn "Please define full_grid(::$(nonparametric_type(G))" diff --git a/src/RingGrids/scaling.jl b/src/RingGrids/scaling.jl index 36e4cc1cb..33a1ad8c7 100644 --- a/src/RingGrids/scaling.jl +++ b/src/RingGrids/scaling.jl @@ -1,10 +1,10 @@ # 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) +scale_coslat!( A::AbstractGridArray) = _scale_coslat!(A, power=1) +scale_coslat²!( A::AbstractGridArray) = _scale_coslat!(A, power=2) +scale_coslat⁻¹!(A::AbstractGridArray) = _scale_coslat!(A, power=-1) +scale_coslat⁻²!(A::AbstractGridArray) = _scale_coslat!(A, power=-2) -function _scale_coslat!(A::Grid; power=1) where {Grid<:AbstractGrid} +function _scale_coslat!(A::Grid; power=1) where {Grid<:AbstractGridArray} coslat = sin.(get_colat(Grid, A.nlat_half)) # sin(colat) = cos(lat) coslat .^= power return _scale_lat!(A, coslat) @@ -13,17 +13,13 @@ 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 +function _scale_lat!(A::AbstractGridArray{NF}, v::AbstractVector) where NF @boundscheck get_nlat(A) == length(v) || throw(BoundsError) - - rings = eachring(A) - - @inbounds for (j, ring) in enumerate(rings) + @inbounds for (j, ring) in enumerate(eachring(A)) vj = convert(NF, v[j]) for ij in ring A[ij] *= vj end end - return A 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) From 28c71349b298c1a23edc769fe9b6641ad3f7497f Mon Sep 17 00:00:00 2001 From: Milan Date: Sun, 14 Apr 2024 17:03:13 -0400 Subject: [PATCH 03/35] scaling for AbstractGridArray --- src/RingGrids/scaling.jl | 45 +++++++++++++++++++++++++--------------- 1 file changed, 28 insertions(+), 17 deletions(-) diff --git a/src/RingGrids/scaling.jl b/src/RingGrids/scaling.jl index 33a1ad8c7..ba6f63a8a 100644 --- a/src/RingGrids/scaling.jl +++ b/src/RingGrids/scaling.jl @@ -1,25 +1,36 @@ -# alias functions to scale the latitude of any gridded map A -scale_coslat!( A::AbstractGridArray) = _scale_coslat!(A, power=1) -scale_coslat²!( A::AbstractGridArray) = _scale_coslat!(A, power=2) -scale_coslat⁻¹!(A::AbstractGridArray) = _scale_coslat!(A, power=-1) -scale_coslat⁻²!(A::AbstractGridArray) = _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<:AbstractGridArray} - 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::AbstractGridArray{NF}, v::AbstractVector) where NF - @boundscheck get_nlat(A) == length(v) || throw(BoundsError) - @inbounds for (j, ring) in enumerate(eachring(A)) - vj = convert(NF, v[j]) - for ij in ring - A[ij] *= vj +function _scale_lat!(grid::AbstractGridArray{T}, v::AbstractVector) where T + @boundscheck get_nlat(grid) == length(v) || throw(BoundsError) + + @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 From 14453c767ec2bc1ce89404cb07175dbe8cba3d6c Mon Sep 17 00:00:00 2001 From: Milan Date: Sun, 14 Apr 2024 17:03:52 -0400 Subject: [PATCH 04/35] similar, eachgrid for AbstractGridArray --- src/RingGrids/full_grids.jl | 80 ++++++++++++++++++------------------- src/RingGrids/general.jl | 75 ++++++++++++++++++++++++++-------- src/RingGrids/similar.jl | 24 ----------- 3 files changed, 96 insertions(+), 83 deletions(-) delete mode 100644 src/RingGrids/similar.jl diff --git a/src/RingGrids/full_grids.jl b/src/RingGrids/full_grids.jl index 561dfcf4c..67320a39d 100644 --- a/src/RingGrids/full_grids.jl +++ b/src/RingGrids/full_grids.jl @@ -10,52 +10,19 @@ const AbstractFullGrid{T} = AbstractFullGridArray{T, 1, Vector{T}} full_grid(G::Type{<:AbstractFullGridArray}) = G - -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) = +## 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)) -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)) -end - -function get_lond(Grid::Type{<:AbstractFullGrid}, nlat_half::Integer) - lon = get_lon(Grid, nlat_half) - lon .*= 360/2π # convert to lond in-place - return lon # = lond -end - +## CONVERSION # 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) - - 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 - colats = zeros(npoints) # preallocate - lons = zeros(npoints) - - for j in 1:nlat # populate preallocated colats, lons - for i in 1:nlon - ij = i + (j-1)*nlon # continuous index ij - colats[ij] = colat[j] - lons[ij] = lon[i] - end - end - - return colats, lons -end +(Grid::Type{<:AbstractFullGrid})(M::AbstractMatrix) = Grid(vec(M)) +Base.Array(grid::AbstractFullGridArray) = Array(reshape(grid.data, :, get_nlat(grid), size(grid.data)[2:end]...)) +## INDEXING function each_index_in_ring( Grid::Type{<:AbstractFullGridArray}, # function for full grids j::Integer, # ring index north to south @@ -83,4 +50,33 @@ function each_index_in_ring!( index_end += nlon # only calculate last index per ring rings[j] = index_1st:index_end # write UnitRange to rings vector end +end + +## 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 + +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 + + npoints = get_npoints(Grid, nlat_half) # total number of grid points + colats = zeros(npoints) # preallocate + lons = zeros(npoints) + + for j in 1:nlat # populate preallocated colats, lons + for i in 1:nlon + ij = i + (j-1)*nlon # continuous index ij + colats[ij] = colat[j] + lons[ij] = lon[i] + end + end + + return colats, lons end \ No newline at end of file diff --git a/src/RingGrids/general.jl b/src/RingGrids/general.jl index 7db89c2a7..e7b127c60 100644 --- a/src/RingGrids/general.jl +++ b/src/RingGrids/general.jl @@ -1,9 +1,6 @@ abstract type AbstractGridArray{T, N, ArrayType <: AbstractArray{T, N}} <: AbstractArray{T, N} end const AbstractGrid{T} = AbstractGridArray{T, 1, Vector{T}} -# nonparametric_type(G::Type{<:AbstractGridArray}) = @warn "Please define nonparametric_type(::$G)" -# horizontal_grid_type(G::Type{<:AbstractGridArray}) = @warn "Please define horizontal_grid_type(::$G)" - # LENGTH SIZE Base.length(G::AbstractGridArray) = length(G.data) Base.size(G::AbstractGridArray) = size(G.data) @@ -21,6 +18,9 @@ get_npoints(grid::Grid) where {Grid<:AbstractGridArray} = get_npoints(Grid, grid get_npoints(G::Type{<:AbstractGridArray}, nlat_half::Integer, k::Integer...) = prod(k) * get_npoints2D(G, nlat_half) get_npoints2D(grid::Grid) where {Grid<:AbstractGridArray} = get_npoints2D(Grid, grid.nlat_half) +# 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 @inline Base.getindex(G::AbstractGridArray, ijk::Integer...) = getindex(G.data,ijk...) @inline Base.getindex(G::AbstractGridArray, r::AbstractRange, k::Integer...) = getindex(G.data, r, k...) @@ -109,7 +109,38 @@ for f in (:zeros, :ones, :rand, :randn) end # zero element of an AbstractGridArray instance grid by creating new zero(grid.data) -Base.zero(grid::Grid) where {Grid<:AbstractGridArray} = Grid(zero(grid.data), grid.nlat_half, grid.rings) +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 + +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})( @@ -164,6 +195,7 @@ get_lond(grid::Grid) where {Grid<:AbstractGridArray} = get_lond(Grid, grid.nlat_ get_lon(grid::Grid) where {Grid<:AbstractGridArray} = get_lon(Grid, grid.nlat_half) get_colat(grid::Grid) where {Grid<:AbstractGridArray} = get_colat(Grid, grid.nlat_half) 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) @@ -177,7 +209,28 @@ end get_lon(::Type{<:AbstractGridArray}, nlat_half::Integer) = Float64[] get_lond(::Type{<:AbstractGridArray}, nlat_half::Integer) = Float64[] +""" +$(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{<:AbstractGridArray}, 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 + ## 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` @@ -187,7 +240,7 @@ and then each grid point per ring. To be used like for ring in rings for ij in ring grid[ij]""" -eachring(grid::AbstractGridArray) = grid.rings +@inline eachring(grid::AbstractGridArray) = grid.rings function eachring(Grid::Type{<:AbstractGridArray}, nlat_half::Integer) rings = Vector{UnitRange{Int}}(undef, get_nlat(Grid, nlat_half)) @@ -245,18 +298,6 @@ function whichring(ij::Integer, rings::Vector{UnitRange{Int}}) 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{<:AbstractGridArray}, 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<:AbstractGridArray} = get_nlon_max(Grid, grid.nlat_half) - ## BROADCASTING # # following https://docs.julialang.org/en/v1/manual/interfaces/#man-interfaces-broadcasting 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 From 88aa2917a70b72505702abfe1865620d376ecdfb Mon Sep 17 00:00:00 2001 From: Milan Date: Sun, 14 Apr 2024 17:04:22 -0400 Subject: [PATCH 05/35] Other full gridarrays --- src/RingGrids/RingGrids.jl | 26 ++++- src/RingGrids/grids/full_clenshaw.jl | 55 +++++----- src/RingGrids/grids/full_healpix.jl | 54 ++++++---- src/RingGrids/grids/full_octahealpix.jl | 53 +++++---- src/RingGrids/grids/octahedral_clenshaw.jl | 120 --------------------- src/RingGrids/grids/octahedral_gaussian.jl | 118 ++++++++++++++++++++ 6 files changed, 232 insertions(+), 194 deletions(-) diff --git a/src/RingGrids/RingGrids.jl b/src/RingGrids/RingGrids.jl index 967073e32..9749db33a 100644 --- a/src/RingGrids/RingGrids.jl +++ b/src/RingGrids/RingGrids.jl @@ -10,17 +10,31 @@ import FastGaussQuadrature import LinearAlgebra # GRIDS +export AbstractGridArray, + AbstractFullGridArray, + AbstractReducedGridArray + export AbstractGrid, AbstractFullGrid, AbstractOctahedralGrid, AbstractHEALPixGrid, AbstractOctaHEALPixGrid +export FullGaussianArray, + FullClenshawArray, + FullHEALPixArray, + FullOctaHEALPixArray + export FullGaussianGrid, FullClenshawGrid, FullHEALPixGrid, FullOctaHEALPixGrid +export OctahedralGaussianArray, + OctahedralClenshawArray, + HEALPixArray, + OctaHEALPixArray + export OctahedralGaussianGrid, OctahedralClenshawGrid, HEALPixGrid, @@ -52,7 +66,11 @@ export grids_match, export scale_coslat!, scale_coslat²!, scale_coslat⁻¹!, - scale_coslat⁻²! + scale_coslat⁻²!, + scale_coslat, + scale_coslat², + scale_coslat⁻¹, + scale_coslat⁻² # INTERPOLATION export AbstractInterpolator, @@ -77,9 +95,9 @@ 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") +include("grids/full_clenshaw.jl") +include("grids/full_healpix.jl") +include("grids/full_octahealpix.jl") # REDUCED GRIDS # include("grids/octahedral_gaussian.jl") diff --git a/src/RingGrids/grids/full_clenshaw.jl b/src/RingGrids/grids/full_clenshaw.jl index 68c96db49..6d9fce596 100644 --- a/src/RingGrids/grids/full_clenshaw.jl +++ b/src/RingGrids/grids/full_clenshaw.jl @@ -1,3 +1,33 @@ +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 + +# 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) + + """ G = FullClenshawGrid{T} @@ -5,27 +35,4 @@ A FullClenshawGrid is a regular latitude-longitude grid with an odd number of `n 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 \ No newline at end of file +FullClenshawGrid \ No newline at end of file diff --git a/src/RingGrids/grids/full_healpix.jl b/src/RingGrids/grids/full_healpix.jl index 2ea857c47..045f1ee43 100644 --- a/src/RingGrids/grids/full_healpix.jl +++ b/src/RingGrids/grids/full_healpix.jl @@ -1,3 +1,33 @@ +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 + +# 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 +get_quadrature_weights(::Type{<:FullHEALPixArray}, nlat_half::Integer) = healpix_weights(nlat_half) + + """ G = FullHEALPixGrid{T} @@ -5,26 +35,4 @@ A full HEALPix grid is a regular latitude-longitude grid that uses `nlat` latitu 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) - +FullHEALPixGrid \ No newline at end of file diff --git a/src/RingGrids/grids/full_octahealpix.jl b/src/RingGrids/grids/full_octahealpix.jl index 76348dee4..52296056f 100644 --- a/src/RingGrids/grids/full_octahealpix.jl +++ b/src/RingGrids/grids/full_octahealpix.jl @@ -1,3 +1,32 @@ +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 + +# 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 +get_quadrature_weights(::Type{<:FullOctaHEALPixArray}, nlat_half::Integer) = healpix_weights(nlat_half) + """ G = FullOctaHEALPixGrid{T} @@ -5,26 +34,4 @@ A full OctaHEALPix grid is a regular latitude-longitude grid that uses `nlat` Oc 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 +FullOctaHEALPixGrid \ No newline at end of file diff --git a/src/RingGrids/grids/octahedral_clenshaw.jl b/src/RingGrids/grids/octahedral_clenshaw.jl index 40149e177..281f115f5 100644 --- a/src/RingGrids/grids/octahedral_clenshaw.jl +++ b/src/RingGrids/grids/octahedral_clenshaw.jl @@ -1,123 +1,3 @@ -""" - 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} diff --git a/src/RingGrids/grids/octahedral_gaussian.jl b/src/RingGrids/grids/octahedral_gaussian.jl index e69de29bb..9ee329813 100644 --- a/src/RingGrids/grids/octahedral_gaussian.jl +++ b/src/RingGrids/grids/octahedral_gaussian.jl @@ -0,0 +1,118 @@ +""" + 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)) + +""" + 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 \ No newline at end of file From 52c9d273a2fc4e45614edb28a117b45b5fad8fea Mon Sep 17 00:00:00 2001 From: Milan Date: Tue, 16 Apr 2024 15:31:28 -0400 Subject: [PATCH 06/35] OctahedralGaussian and ClenshawArrays --- src/RingGrids/RingGrids.jl | 9 +- src/RingGrids/general.jl | 1 + src/RingGrids/grids/octahedral_clenshaw.jl | 137 ++++++++++++++------ src/RingGrids/grids/octahedral_gaussian.jl | 139 ++++++++++----------- src/RingGrids/reduced_grids.jl | 26 ++++ 5 files changed, 197 insertions(+), 115 deletions(-) diff --git a/src/RingGrids/RingGrids.jl b/src/RingGrids/RingGrids.jl index 9749db33a..0d78bb4e9 100644 --- a/src/RingGrids/RingGrids.jl +++ b/src/RingGrids/RingGrids.jl @@ -90,7 +90,7 @@ include("utility_functions.jl") # GENERAL include("general.jl") include("full_grids.jl") -# include("reduced_grids.jl") +include("reduced_grids.jl") include("scaling.jl") # FULL GRIDS @@ -100,15 +100,14 @@ include("grids/full_healpix.jl") include("grids/full_octahealpix.jl") # REDUCED GRIDS -# include("grids/octahedral_gaussian.jl") -# include("grids/octahedral_clenshaw.jl") +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") -# include("similar.jl") +include("interpolation.jl") # OUTPUT include("show.jl") diff --git a/src/RingGrids/general.jl b/src/RingGrids/general.jl index e7b127c60..591507ec1 100644 --- a/src/RingGrids/general.jl +++ b/src/RingGrids/general.jl @@ -132,6 +132,7 @@ function Base.similar( return nonparametric_type(Grid)(similar_data, nlat_half) end +# data with new type T and new size function Base.similar( grid::Grid, ::Type{Tnew}, diff --git a/src/RingGrids/grids/octahedral_clenshaw.jl b/src/RingGrids/grids/octahedral_clenshaw.jl index 281f115f5..e5d30e778 100644 --- a/src/RingGrids/grids/octahedral_clenshaw.jl +++ b/src/RingGrids/grids/octahedral_clenshaw.jl @@ -1,40 +1,64 @@ -""" - G = OctahedralClenshawGrid{T} +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 -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}.") +# TYPES +const OctahedralClenshawGrid{T} = OctahedralClenshawArray{T, 1, Vector{T}} +nonparametric_type(::Type{<:OctahedralClenshawArray}) = OctahedralClenshawArray +horizontal_grid_type(::Type{<:OctahedralClenshawArray}) = OctahedralClenshawGrid +full_grid(::Type{<:OctahedralClenshawArray}) = FullClenshawArray + +# SIZE +nlat_odd(::Type{<:OctahedralClenshawArray}) = true +npoints_pole(::Type{<:OctahedralClenshawArray}) = 0 +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 -nonparametric_type(::Type{<:OctahedralClenshawGrid}) = OctahedralClenshawGrid +function get_nlat_half(::Type{<:OctahedralClenshawArray}, npoints2D::Integer) + m, o = npoints_added_per_ring(OctahedralClenshawArray), npoints_pole(OctahedraClenshawArray) + return round(Int, -o/m + sqrt(((o/m)^2 + (o+npoints2D)/m))) +end -# 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...) +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 -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(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(OctahedraClenshawArray) + m != 4 && @warn "This algorithm has not been generalised for m!=4." + N = (o + 4n)÷2 + return (N, N) +end -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...) +## 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) + +## CONVERSION """ Matrix!(M::AbstractMatrix, G::OctahedralClenshawGrid; @@ -60,7 +84,11 @@ 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 @@ -95,9 +123,9 @@ function Matrix!( MGs::Tuple{AbstractMatrix{T}, OctahedralClenshawGrid}...; 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) + 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]) @@ -107,12 +135,47 @@ function Matrix!( MGs::Tuple{AbstractMatrix{T}, OctahedralClenshawGrid}...; 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 + 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 + return Tuple(Mi for (Mi, Gi) in MGs) +end + +## INDEXING +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 + +""" + 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.""" \ No newline at end of file diff --git a/src/RingGrids/grids/octahedral_gaussian.jl b/src/RingGrids/grids/octahedral_gaussian.jl index 9ee329813..2cf7428e5 100644 --- a/src/RingGrids/grids/octahedral_gaussian.jl +++ b/src/RingGrids/grids/octahedral_gaussian.jl @@ -1,94 +1,74 @@ -""" - 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}.") +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 -nonparametric_type(::Type{<:OctahedralGaussianGrid}) = OctahedralGaussianGrid +# TYPES +const OctahedralGaussianGrid{T} = OctahedralGaussianArray{T, 1, Vector{T}} +nonparametric_type(::Type{<:OctahedralGaussianArray}) = OctahedralGaussianArray +horizontal_grid_type(::Type{<:OctahedralGaussianArray}) = OctahedralGaussianGrid +full_grid(::Type{<:OctahedralGaussianArray}) = FullGaussianArray -# 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 +# SIZE +nlat_odd(::Type{<:OctahedralGaussianArray}) = false +npoints_pole(::Type{<:OctahedralGaussianArray}) = 16 +npoints_added_per_ring(::Type{<:OctahedralGaussianArray}) = 4 -# 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)) - -""" - 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 +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 -get_nlon_max(::Type{<:AbstractOctahedralGrid}, nlat_half::Integer) = nlon_octahedral(nlat_half) +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{<:AbstractOctahedralGrid}, nlat_half::Integer, j::Integer) +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 nlon_octahedral(j) + return o + m*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) +# maybe define at some point for Matrix(::OctahedralGaussianGrid) +# matrix_size(G::OctahedralGaussianGrid) = (2*(4+G.nlat_half), 2*(4+G.nlat_half+1)) - colats[ij:ij+nlon-1] .= colat[j] - lons[ij:ij+nlon-1] .= lon - - ij += nlon - end - - return colats, lons +## 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 -function each_index_in_ring(Grid::Type{<:AbstractOctahedralGrid}, +## QUADRATURE +get_quadrature_weights(::Type{<:OctahedralGaussianArray}, nlat_half::Integer) = gaussian_weights(nlat_half) + +## INDEXING +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_npoints(Grid, nlat_half) + 1 # number of grid points + 1 + 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 @@ -96,23 +76,36 @@ function each_index_in_ring(Grid::Type{<:AbstractOctahedralGrid}, end function each_index_in_ring!( rings::Vector{<:UnitRange{<:Integer}}, - Grid::Type{<:AbstractOctahedralGrid}, + 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 += 16 + 4j # add number of grid points per ring + 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_mir) in zip( nlat_half+1:nlat, # South only - nlat-nlat_half:-1:1) # reverse index + @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 += 16 + 4j_mir # add number of grid points per ring + 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 +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. +# """ +# # OctahedralGaussianGrid \ No newline at end of file diff --git a/src/RingGrids/reduced_grids.jl b/src/RingGrids/reduced_grids.jl index 97b0da886..766deecfe 100644 --- a/src/RingGrids/reduced_grids.jl +++ b/src/RingGrids/reduced_grids.jl @@ -2,3 +2,29 @@ abstract type AbstractReducedGridArray{T, N, ArrayType <: AbstractArray{T, N}} < const AbstractReducedGrid{T} = AbstractReducedGridArray{T, 1, Vector{T}} full_grid(G::Type{<:AbstractReducedGridArray}) = @warn "Please define full_grid(::$(nonparametric_type(G))" + +# 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 From b4386253c25cab8a979ef3d25a91fadf85d69151 Mon Sep 17 00:00:00 2001 From: Milan Date: Tue, 16 Apr 2024 19:23:41 -0400 Subject: [PATCH 07/35] Tests pass except broadcasting --- docs/src/output.md | 6 +- ext/SpeedyWeatherMakieExt.jl | 2 +- src/RingGrids/RingGrids.jl | 6 +- src/RingGrids/full_grids.jl | 4 +- src/RingGrids/general.jl | 10 +- src/RingGrids/grids/healpix.jl | 134 +++++++++------------ src/RingGrids/grids/octahealpix.jl | 132 ++++++++------------ src/RingGrids/grids/octahedral_clenshaw.jl | 55 ++++++++- src/RingGrids/grids/octahedral_gaussian.jl | 2 +- src/RingGrids/quadrature_weights.jl | 21 ++-- src/RingGrids/reduced_grids.jl | 2 - src/RingGrids/show.jl | 2 +- src/output/output.jl | 4 +- test/grids.jl | 36 +++--- 14 files changed, 216 insertions(+), 200 deletions(-) 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/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 0d78bb4e9..9c49d0921 100644 --- a/src/RingGrids/RingGrids.jl +++ b/src/RingGrids/RingGrids.jl @@ -102,11 +102,11 @@ 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") +include("grids/healpix.jl") +include("grids/octahealpix.jl") # INTEGRATION AND INTERPOLATION -# include("quadrature_weights.jl") +include("quadrature_weights.jl") include("interpolation.jl") # OUTPUT diff --git a/src/RingGrids/full_grids.jl b/src/RingGrids/full_grids.jl index 67320a39d..c63d9ccb0 100644 --- a/src/RingGrids/full_grids.jl +++ b/src/RingGrids/full_grids.jl @@ -7,8 +7,7 @@ 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.""" const AbstractFullGrid{T} = AbstractFullGridArray{T, 1, Vector{T}} - -full_grid(G::Type{<:AbstractFullGridArray}) = G +full_grid_type(Grid::Type{<:AbstractFullGridArray}) = nonparametric_type(Grid) ## SIZE get_nlon_max(Grid::Type{<:AbstractFullGridArray}, nlat_half::Integer) = get_nlon(Grid, nlat_half) @@ -21,6 +20,7 @@ matrix_size(Grid::Type{<:AbstractFullGridArray}, nlat_half::Integer) = # 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 function each_index_in_ring( diff --git a/src/RingGrids/general.jl b/src/RingGrids/general.jl index 591507ec1..42a538578 100644 --- a/src/RingGrids/general.jl +++ b/src/RingGrids/general.jl @@ -1,7 +1,15 @@ abstract type AbstractGridArray{T, N, ArrayType <: AbstractArray{T, N}} <: AbstractArray{T, N} end + +# Horizontal 2D grids with N=1 const AbstractGrid{T} = AbstractGridArray{T, 1, Vector{T}} -# LENGTH SIZE +## TYPES +nonparametric_type(grid::AbstractGridArray) = nonparametric_type(typeof(grid)) +full_array_type(grid::AbstractGridArray) = full_array_type(typeof(grid)) +full_grid_type(grid::AbstractGridArray) = horizontal_grid_type(full_array_type(grid)) +horizontal_grid_type(grid::AbstractGridArray) = horizontal_grid_type(typeof(grid)) + +## SIZE Base.length(G::AbstractGridArray) = length(G.data) Base.size(G::AbstractGridArray) = size(G.data) Base.sizeof(G::AbstractGridArray) = sizeof(G.data) diff --git a/src/RingGrids/grids/healpix.jl b/src/RingGrids/grids/healpix.jl index 38846ada1..6bd47058c 100644 --- a/src/RingGrids/grids/healpix.jl +++ b/src/RingGrids/grids/healpix.jl @@ -1,54 +1,63 @@ -""" - 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) +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 -function get_nlon_per_ring(G::Type{<:AbstractHEALPixGrid}, nlat_half::Integer, j::Integer) - nlat = get_nlat(G, nlat_half) +## 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 + +## 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)) +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 nlon_healpix(nlat_half, j) + return min(4j, 2nlat_half, 8nlat_half-4j) 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_healpix(nlat_half::Integer) = nlat_half÷2 +function matrix_size(::Type{<:HEALPixArray}, nlat_half::Integer) 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) + return (5nside, 5nside) +end - # 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] +## COORDINATES +function get_colat(::Type{<:HEALPixArray}, nlat_half::Integer) + nlat_half == 0 && return Float64[] - colats[ij:ij+nlon-1] .= colat[j] - lons[ij:ij+nlon-1] .= lon + nlat = get_nlat(HEALPixArray, 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 - ij += nlon - end + 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) - return colats, lons + # 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 -function each_index_in_ring(::Type{<:AbstractHEALPixGrid}, # function for HEALPix grids +## 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) @@ -73,7 +82,7 @@ function each_index_in_ring(::Type{<:AbstractHEALPixGrid}, # function for HEALP end function each_index_in_ring!( rings::Vector{<:UnitRange{<:Integer}}, - Grid::Type{<:AbstractHEALPixGrid}, + Grid::Type{<:HEALPixArray}, nlat_half::Integer) # resolution param nlat = length(rings) @@ -90,7 +99,7 @@ function each_index_in_ring!( rings::Vector{<:UnitRange{<:Integer}}, end # Equatorial belt - nlon_max = get_nlon_max(Grid, nlat_half) # number of grid points on 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 @@ -98,11 +107,11 @@ function each_index_in_ring!( rings::Vector{<:UnitRange{<:Integer}}, end # South polar cap - @inbounds for (j, j_mir) in zip( 3nside+1:nlat, # South only - nside-1:-1:1) # mirror index + @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_mir # add number of grid points per ring + index_end += 4j_mirrored # add number of grid points per ring rings[j] = index_1st:index_end # turn into UnitRange end end @@ -114,37 +123,4 @@ A HEALPix grid with 12 faces, each `nside`x`nside` grid points, each covering th 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 +HEALPixGrid \ No newline at end of file diff --git a/src/RingGrids/grids/octahealpix.jl b/src/RingGrids/grids/octahealpix.jl index a2c5f344f..6287cd736 100644 --- a/src/RingGrids/grids/octahealpix.jl +++ b/src/RingGrids/grids/octahealpix.jl @@ -1,47 +1,51 @@ -""" - 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) +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 - -function get_colatlons(Grid::Type{<:AbstractOctaHEALPixGrid}, nlat_half::Integer) +## 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 + +## 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)) +function get_nlon_per_ring(Grid::Type{<:OctaHEALPixArray}, nlat_half::Integer, j::Integer) nlat = get_nlat(Grid, nlat_half) - npoints = get_npoints(Grid, 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) - lon = collect(π/nlon:2π/nlon:2π) + @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 + 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 each_index_in_ring(::Type{<:AbstractOctaHEALPixGrid}, # function for OctaHEALPix grids +function get_lon_per_ring(Grid::Type{<:OctaHEALPixArray}, nlat_half::Integer, j::Integer) + nlon = get_nlon_per_ring(Grid, nlat_half, j) + return collect(π/nlon:2π/nlon:2π) +end + +## 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 +63,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 +75,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 +84,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...) """ @@ -203,4 +167,12 @@ function Matrix!( MGs::Tuple{AbstractMatrix{T}, OctaHEALPixGrid}...; ntuples == 1 && return M return Tuple(Mi for (Mi, Gi) in MGs) -end \ No newline at end of file +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.""" +OctaHEALPixGrid \ No newline at end of file diff --git a/src/RingGrids/grids/octahedral_clenshaw.jl b/src/RingGrids/grids/octahedral_clenshaw.jl index e5d30e778..7388470b3 100644 --- a/src/RingGrids/grids/octahedral_clenshaw.jl +++ b/src/RingGrids/grids/octahedral_clenshaw.jl @@ -13,11 +13,11 @@ end const OctahedralClenshawGrid{T} = OctahedralClenshawArray{T, 1, Vector{T}} nonparametric_type(::Type{<:OctahedralClenshawArray}) = OctahedralClenshawArray horizontal_grid_type(::Type{<:OctahedralClenshawArray}) = OctahedralClenshawGrid -full_grid(::Type{<:OctahedralClenshawArray}) = FullClenshawArray +full_array_type(::Type{<:OctahedralClenshawArray}) = FullClenshawArray # SIZE nlat_odd(::Type{<:OctahedralClenshawArray}) = true -npoints_pole(::Type{<:OctahedralClenshawArray}) = 0 +npoints_pole(::Type{<:OctahedralClenshawArray}) = 16 npoints_added_per_ring(::Type{<:OctahedralClenshawArray}) = 4 function get_npoints2D(::Type{<:OctahedralClenshawArray}, nlat_half::Integer) @@ -26,7 +26,7 @@ function get_npoints2D(::Type{<:OctahedralClenshawArray}, nlat_half::Integer) end function get_nlat_half(::Type{<:OctahedralClenshawArray}, npoints2D::Integer) - m, o = npoints_added_per_ring(OctahedralClenshawArray), npoints_pole(OctahedraClenshawArray) + m, o = npoints_added_per_ring(OctahedralClenshawArray), npoints_pole(OctahedralClenshawArray) return round(Int, -o/m + sqrt(((o/m)^2 + (o+npoints2D)/m))) end @@ -40,7 +40,7 @@ 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(OctahedraClenshawArray) + m, o = npoints_added_per_ring(OctahedralClenshawArray), npoints_pole(OctahedralClenshawArray) m != 4 && @warn "This algorithm has not been generalised for m!=4." N = (o + 4n)÷2 return (N, N) @@ -58,6 +58,53 @@ 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 """ Matrix!(M::AbstractMatrix, diff --git a/src/RingGrids/grids/octahedral_gaussian.jl b/src/RingGrids/grids/octahedral_gaussian.jl index 2cf7428e5..7a33b36d8 100644 --- a/src/RingGrids/grids/octahedral_gaussian.jl +++ b/src/RingGrids/grids/octahedral_gaussian.jl @@ -13,7 +13,7 @@ end const OctahedralGaussianGrid{T} = OctahedralGaussianArray{T, 1, Vector{T}} nonparametric_type(::Type{<:OctahedralGaussianArray}) = OctahedralGaussianArray horizontal_grid_type(::Type{<:OctahedralGaussianArray}) = OctahedralGaussianGrid -full_grid(::Type{<:OctahedralGaussianArray}) = FullGaussianArray +full_array_type(::Type{<:OctahedralGaussianArray}) = FullGaussianArray # SIZE nlat_odd(::Type{<:OctahedralGaussianArray}) = false diff --git a/src/RingGrids/quadrature_weights.jl b/src/RingGrids/quadrature_weights.jl index 3a51673dc..983728f2b 100644 --- a/src/RingGrids/quadrature_weights.jl +++ b/src/RingGrids/quadrature_weights.jl @@ -3,19 +3,26 @@ # 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] 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]] 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)) +function equal_area_weights(Grid::Type{<:AbstractGridArray}, nlat_half::Integer) + weights = zeros(nlat_half) + for j in 1:nlat_half + nlon = get_nlon_per_ring(Grid, nlat_half, j) + weights[j] = 2nlon/get_npoints2D(Grid, nlat_half) + end + return weights +end + +healpix_weights(nlat_half::Integer) = equal_area_weights(HEALPixArray, nlat_half) +octahealpix_weights(nlat_half::Integer) = equal_area_weights(OctaHEALPixArray, nlat_half) # 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) = +get_solid_angles(Grid::Type{<:Union{HEALPixArray, OctaHEALPixArray}}, nlat_half::Integer) = 4π/get_npoints(Grid, nlat_half)*ones(nlat_half) diff --git a/src/RingGrids/reduced_grids.jl b/src/RingGrids/reduced_grids.jl index 766deecfe..c41571b6e 100644 --- a/src/RingGrids/reduced_grids.jl +++ b/src/RingGrids/reduced_grids.jl @@ -1,8 +1,6 @@ abstract type AbstractReducedGridArray{T, N, ArrayType <: AbstractArray{T, N}} <: AbstractGridArray{T, N, ArrayType} end const AbstractReducedGrid{T} = AbstractReducedGridArray{T, 1, Vector{T}} -full_grid(G::Type{<:AbstractReducedGridArray}) = @warn "Please define full_grid(::$(nonparametric_type(G))" - # 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) diff --git a/src/RingGrids/show.jl b/src/RingGrids/show.jl index 88304d088..9d02e6795 100644 --- a/src/RingGrids/show.jl +++ b/src/RingGrids/show.jl @@ -7,7 +7,7 @@ function Base.array_summary(io::IO, grid::AbstractGrid, inds::Tuple{Vararg{Base. 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/output/output.jl b/src/output/output.jl index c946c2d51..4dc815619 100644 --- a/src/output/output.jl +++ b/src/output/output.jl @@ -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 diff --git a/test/grids.jl b/test/grids.jl index 51a0f15f4..b1d04a707 100644 --- a/test/grids.jl +++ b/test/grids.jl @@ -2,20 +2,20 @@ 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 +204,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 @@ -230,6 +230,14 @@ 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 + grid = zeros(G, n) + @test typeof(grid + grid) isa G + @test typeof(grid - grid) isa G + @test typeof(grid .* grid) isa G + @test typeof(grid ./ grid) isa G + @test typeof(2grid) isa G + # 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) From 0ba0612bb854fcdc61710290d9f7cada6acd2ff8 Mon Sep 17 00:00:00 2001 From: Milan Date: Wed, 17 Apr 2024 10:25:50 -0400 Subject: [PATCH 08/35] N-dimensional broadcasting for RingGrids --- src/RingGrids/general.jl | 34 +++++++++++++++------------------- src/RingGrids/show.jl | 2 +- test/grids.jl | 10 +++++----- 3 files changed, 21 insertions(+), 25 deletions(-) diff --git a/src/RingGrids/general.jl b/src/RingGrids/general.jl index 42a538578..c6da5562d 100644 --- a/src/RingGrids/general.jl +++ b/src/RingGrids/general.jl @@ -308,26 +308,22 @@ function whichring(ij::Integer, rings::Vector{UnitRange{Int}}) end ## BROADCASTING +# following https://docs.julialang.org/en/v1/manual/interfaces/#man-interfaces-broadcasting +import Base.Broadcast: BroadcastStyle, Broadcasted -# # 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 AbstractGridArrayStyle{N, Grid} <: Broadcast.AbstractArrayStyle{N} end -# # {1} as grids are <:AbstractVector, Grid here is the non-parameteric Grid type! -# struct AbstractGridArrayStyle{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<:AbstractGridArray{T, N, ArrayType}} where {T, N, ArrayType} = + AbstractGridArrayStyle{N, nonparametric_type(Grid)}() -# # 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} = -# AbstractGridArrayStyle{nonparametric_type(Grid)}() - -# # allocation for broadcasting, create a new Grid with undef of type/number format T -# function Base.similar(bc::Broadcasted{AbstractGridArrayStyle{Grid}}, ::Type{T}) where {Grid, T} -# Grid(Vector{T}(undef,length(bc))) -# end +# 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{Grid}(::Val{0}) where Grid = AbstractGridArrayStyle{Grid}() -# AbstractGridArrayStyle{Grid}(::Val{1}) where Grid = AbstractGridArrayStyle{Grid}() -# AbstractGridArrayStyle{Grid}(::Val{2}) where Grid = AbstractGridArrayStyle{Grid}() -# AbstractGridArrayStyle{Grid}(::Val{3}) where Grid = AbstractGridArrayStyle{Grid}() \ No newline at end of file +# ::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}() \ No newline at end of file diff --git a/src/RingGrids/show.jl b/src/RingGrids/show.jl index 9d02e6795..51ea94095 100644 --- a/src/RingGrids/show.jl +++ b/src/RingGrids/show.jl @@ -1,7 +1,7 @@ # 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 diff --git a/test/grids.jl b/test/grids.jl index b1d04a707..1ef45cc4e 100644 --- a/test/grids.jl +++ b/test/grids.jl @@ -232,11 +232,11 @@ end # don't promote to Array grid = zeros(G, n) - @test typeof(grid + grid) isa G - @test typeof(grid - grid) isa G - @test typeof(grid .* grid) isa G - @test typeof(grid ./ grid) isa G - @test typeof(2grid) 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 # promote types, Grid{Float16} -> Grid{Float64} etc @test all(ones(G{Float16}, n)*2.0 .=== 2.0) From a51164ca03ad915513666f14fec3cc6eef3d8bf6 Mon Sep 17 00:00:00 2001 From: Milan Date: Wed, 17 Apr 2024 11:03:04 -0400 Subject: [PATCH 09/35] only two getindex methods --- src/RingGrids/general.jl | 11 +++----- test/grids.jl | 59 ++++++++++++++++++++++++++++++---------- 2 files changed, 49 insertions(+), 21 deletions(-) diff --git a/src/RingGrids/general.jl b/src/RingGrids/general.jl index c6da5562d..a0c75a453 100644 --- a/src/RingGrids/general.jl +++ b/src/RingGrids/general.jl @@ -30,18 +30,15 @@ get_npoints2D(grid::Grid) where {Grid<:AbstractGridArray} = get_npoints2D(Grid, matrix_size(grid::Grid) where {Grid<:AbstractGridArray} = matrix_size(Grid, grid.nlat_half) ## INDEXING -@inline Base.getindex(G::AbstractGridArray, ijk::Integer...) = getindex(G.data,ijk...) -@inline Base.getindex(G::AbstractGridArray, r::AbstractRange, k::Integer...) = getindex(G.data, r, k...) -@inline Base.getindex(G::AbstractGridArray, ij::Integer, k::AbstractRange) = getindex(G.data, ij, k) -@inline Base.getindex(G::AbstractGridArray, I::CartesianIndex) = getindex(G, Tuple(I)...) +@inline Base.getindex(G::AbstractGridArray, ijk...) = getindex(G.data, ijk...) @inline function Base.getindex( G::GridArray, col::Colon, - k..., -) where {GridArray<:AbstractGridArray{T, N, ArrayType}} where {T, N, ArrayType} + k::Integer..., +) where {GridArray<:AbstractGridArray} GridArray_ = nonparametric_type(GridArray) - return GridArray_{T, 1}(getindex(G.data, col, k...), G.nlat_half) + return GridArray_(getindex(G.data, col, k...), G.nlat_half, G.rings) end @inline Base.setindex!(G::AbstractGridArray, x, ijk::Integer...) = diff --git a/test/grids.jl b/test/grids.jl index 1ef45cc4e..7ca6362d8 100644 --- a/test/grids.jl +++ b/test/grids.jl @@ -214,14 +214,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) @@ -231,12 +231,14 @@ end @test 2ones(G, n) == ones(G, n) + ones(G, n) # don't promote to Array - grid = zeros(G, n) - @test (grid + grid) isa G - @test (grid - grid) isa G - @test (grid .* grid) isa G - @test (grid ./ grid) isa G - @test 2grid isa G + 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 2grid isa G + end # promote types, Grid{Float16} -> Grid{Float64} etc @test all(ones(G{Float16}, n)*2.0 .=== 2.0) @@ -263,4 +265,33 @@ 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 \ No newline at end of file From 652a6d09ae265d92b84037f31e7e12eaf6b63340 Mon Sep 17 00:00:00 2001 From: Milan Date: Wed, 17 Apr 2024 11:15:22 -0400 Subject: [PATCH 10/35] Loop indexing tested --- src/RingGrids/RingGrids.jl | 29 +++++++++++++++++------------ test/grids.jl | 27 +++++++++++++++++++++++++++ 2 files changed, 44 insertions(+), 12 deletions(-) diff --git a/src/RingGrids/RingGrids.jl b/src/RingGrids/RingGrids.jl index 9c49d0921..09e514177 100644 --- a/src/RingGrids/RingGrids.jl +++ b/src/RingGrids/RingGrids.jl @@ -40,28 +40,33 @@ export OctahedralGaussianGrid, 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²!, diff --git a/test/grids.jl b/test/grids.jl index 7ca6362d8..e123c0531 100644 --- a/test/grids.jl +++ b/test/grids.jl @@ -294,4 +294,31 @@ end @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 \ No newline at end of file From 60214673a00e45be59db4c55b694f06767112641 Mon Sep 17 00:00:00 2001 From: Milan Date: Wed, 17 Apr 2024 11:56:19 -0400 Subject: [PATCH 11/35] Export GridArrays too --- src/RingGrids/general.jl | 3 +-- src/RingGrids/interpolation.jl | 16 ++++++++-------- src/SpeedyWeather.jl | 16 ++++++++-------- 3 files changed, 17 insertions(+), 18 deletions(-) diff --git a/src/RingGrids/general.jl b/src/RingGrids/general.jl index a0c75a453..90752bb34 100644 --- a/src/RingGrids/general.jl +++ b/src/RingGrids/general.jl @@ -258,8 +258,7 @@ 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<:AbstractGridArray} - @inline +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) diff --git a/src/RingGrids/interpolation.jl b/src/RingGrids/interpolation.jl index bacbcdc60..7ebed411e 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,11 +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 + 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) lon_offsets = [londs[ring[1]] for ring in rings]# offset of the first point from 0˚E diff --git a/src/SpeedyWeather.jl b/src/SpeedyWeather.jl index 753f5de12..495ad2e63 100644 --- a/src/SpeedyWeather.jl +++ b/src/SpeedyWeather.jl @@ -43,14 +43,14 @@ using .LowerTriangularMatrices # RingGrids export RingGrids -export FullClenshawGrid, - FullGaussianGrid, - FullHEALPixGrid, - FullOctaHEALPixGrid, - OctahedralGaussianGrid, - OctahedralClenshawGrid, - HEALPixGrid, - OctaHEALPixGrid, +export FullClenshawGrid, FullClenshawArray, + FullGaussianGrid, FullGaussianArray, + FullHEALPixGrid, FullHEALPixArray, + FullOctaHEALPixGrid, FullOctaHEALPixArray, + OctahedralGaussianGrid, OctahedralGaussianArray, + OctahedralClenshawGrid, OctahedralClenshawArray, + HEALPixGrid, HEALPixArray, + OctaHEALPixGrid, OctaHEALPixArray, plot include("RingGrids/RingGrids.jl") From 5971060b5cd12a57df6adc3c2618938a854a72f8 Mon Sep 17 00:00:00 2001 From: Milan Date: Wed, 17 Apr 2024 12:11:09 -0400 Subject: [PATCH 12/35] export also eachgrid, eachring --- src/SpeedyWeather.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/SpeedyWeather.jl b/src/SpeedyWeather.jl index 495ad2e63..05ea83605 100644 --- a/src/SpeedyWeather.jl +++ b/src/SpeedyWeather.jl @@ -51,7 +51,7 @@ export FullClenshawGrid, FullClenshawArray, OctahedralClenshawGrid, OctahedralClenshawArray, HEALPixGrid, HEALPixArray, OctaHEALPixGrid, OctaHEALPixArray, - plot + eachring, eachgrid, plot include("RingGrids/RingGrids.jl") using .RingGrids From 917b1a0dd7cc1c1e48892eb42b21f8678e0a5fa0 Mon Sep 17 00:00:00 2001 From: Milan Date: Wed, 17 Apr 2024 14:40:52 -0400 Subject: [PATCH 13/35] add Adapt, KernelAbstractions, GPUArrays deps, JLArrays extra --- Project.toml | 5 ++++- src/RingGrids/RingGrids.jl | 5 +++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index f89f3cb1b..c952774f1 100644 --- a/Project.toml +++ b/Project.toml @@ -15,6 +15,7 @@ DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341" FLoops = "cc61a311-1640-44b5-9fba-1b764f453329" FastGaussQuadrature = "442a2c76-b920-505d-bb47-c5924d526838" +GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7" GenericFFT = "a8297547-1b15-4a5a-a998-a2ac5f1cef28" JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" @@ -46,6 +47,7 @@ DocStringExtensions = "0.9" FFTW = "1" FLoops = "0.2" FastGaussQuadrature = "0.4, 0.5, 1" +GPUArrays = "10" GenericFFT = "0.1" JLD2 = "0.4" KernelAbstractions = "0.9" @@ -62,7 +64,8 @@ UnicodePlots = "3.3" julia = "1.9" [extras] +JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["Test"] +test = ["Test", "JLArrays"] diff --git a/src/RingGrids/RingGrids.jl b/src/RingGrids/RingGrids.jl index 09e514177..7a6ffcbc7 100644 --- a/src/RingGrids/RingGrids.jl +++ b/src/RingGrids/RingGrids.jl @@ -9,6 +9,11 @@ import Statistics: mean import FastGaussQuadrature import LinearAlgebra +# GPU +import Adapt +import KernelAbstractions +import GPUArrays + # GRIDS export AbstractGridArray, AbstractFullGridArray, From 3ac3b6eae60afd9b8cac243e522daedac1f13558 Mon Sep 17 00:00:00 2001 From: Milan Date: Wed, 17 Apr 2024 14:41:18 -0400 Subject: [PATCH 14/35] GPU broadcasting first draft --- src/RingGrids/general.jl | 30 +++++++++++++-- src/RingGrids/grids/octahedral_clenshaw.jl | 24 ------------ test/grids.jl | 44 ++++++++++++++++++++++ 3 files changed, 71 insertions(+), 27 deletions(-) diff --git a/src/RingGrids/general.jl b/src/RingGrids/general.jl index 90752bb34..b118051a2 100644 --- a/src/RingGrids/general.jl +++ b/src/RingGrids/general.jl @@ -7,6 +7,7 @@ const AbstractGrid{T} = AbstractGridArray{T, 1, Vector{T}} nonparametric_type(grid::AbstractGridArray) = nonparametric_type(typeof(grid)) full_array_type(grid::AbstractGridArray) = full_array_type(typeof(grid)) 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)) horizontal_grid_type(grid::AbstractGridArray) = horizontal_grid_type(typeof(grid)) ## SIZE @@ -305,7 +306,8 @@ end ## BROADCASTING # following https://docs.julialang.org/en/v1/manual/interfaces/#man-interfaces-broadcasting -import Base.Broadcast: BroadcastStyle, Broadcasted +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 @@ -318,8 +320,30 @@ Base.BroadcastStyle(::Type{Grid}) where {Grid<:AbstractGridArray{T, N, ArrayType # 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)...)) + # if isstructurepreserving(bc) || fzeropreserving(bc) + # return Grid(Array{T}(undef, size(bc)...)) + # end + # return similar(convert(Broadcasted{DefaultArrayStyle{ndims(bc)}}, bc), 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}() \ No newline at end of file +AbstractGridArrayStyle{N, Grid}(::Val{M}) where {N, Grid, M} = AbstractGridArrayStyle{N, Grid}() + +## GPU +function Broadcast.BroadcastStyle( + ::Type{Grid}, +) where {Grid <: AbstractGridArray{T, N, ArrayType}} where {T, N, ArrayType <: GPUArrays.AbstractGPUArray} + return Broadcast.BroadcastStyle(ArrayType) +end + +function GPUArrays.backend( + ::Type{Grid}, +) where {Grid <: AbstractGridArray{T, N, ArrayType}} where {T, N, ArrayType <: GPUArrays.AbstractGPUArray} + return GPUArrays.backend(ArrayType) +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/octahedral_clenshaw.jl b/src/RingGrids/grids/octahedral_clenshaw.jl index 7388470b3..593c9a88c 100644 --- a/src/RingGrids/grids/octahedral_clenshaw.jl +++ b/src/RingGrids/grids/octahedral_clenshaw.jl @@ -192,30 +192,6 @@ function Matrix!( MGs::Tuple{AbstractMatrix{T}, OctahedralClenshawGrid}...; return Tuple(Mi for (Mi, Gi) in MGs) end -## INDEXING -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 - """ G = OctahedralClenshawGrid{T} diff --git a/test/grids.jl b/test/grids.jl index e123c0531..1eb1a86f0 100644 --- a/test/grids.jl +++ b/test/grids.jl @@ -321,4 +321,48 @@ end @test all(grid .== 1) end end +end + +using JLArrays + +@testset "AbstractGridArray: GPU (JLArrays)" begin + NF = Float32 + @testset for Grid in ( + FullClenshawArray, + FullGaussianArray, + OctahedralGaussianArray, + OctahedralClenshawArray, + HEALPixArray, + OctaHEALPixArray, + FullHEALPixArray, + FullOctaHEALPixArray, + ) + G_cpu = randn(Grid{NF}, 2, 3, 4) + + # constructors/adapt + G = adapt(JLArray, G_cpu) + G2 = Grid(adapt(JLArray, G_cpu.data)) + # @test G == G2 + # @test all(G .== G2) + + # getindex + # @test G[1, :] isa JLArray + # @test G[:, 1] isa Grid{NF, 1, JLArray{NF}} + # for k in eachgrid(G) + # for (j, ring) in enumerate(eachring(G)) + # @test Array(L[lm,:]) == L_cpu[lm,:] + # end + # end + + # # setindex! + # A_test = JLArray(rand(NF,size(L_cpu,3))) + # L[1,:] = A_test + # @test L[1,:] == A_test + + # # fill + # fill!(L, 2) + # for lm in SpeedyWeather.eachharmonic(L2) + # @test all(L[lm, [Colon() for i=1:length(idims)]...] .== 2) + # end + end end \ No newline at end of file From 963988b5814bbdfc00a05c96ee57ef703726ef04 Mon Sep 17 00:00:00 2001 From: Milan Date: Wed, 17 Apr 2024 18:51:53 -0400 Subject: [PATCH 15/35] using Adapt, JLArrays in ringgrids test --- src/RingGrids/general.jl | 3 +++ test/grids.jl | 5 +++-- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/src/RingGrids/general.jl b/src/RingGrids/general.jl index b118051a2..f57c03ab1 100644 --- a/src/RingGrids/general.jl +++ b/src/RingGrids/general.jl @@ -320,10 +320,13 @@ Base.BroadcastStyle(::Type{Grid}) where {Grid<:AbstractGridArray{T, N, ArrayType # 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} + # this escapes for Array and JLArray # if isstructurepreserving(bc) || fzeropreserving(bc) # return Grid(Array{T}(undef, size(bc)...)) # end # return similar(convert(Broadcasted{DefaultArrayStyle{ndims(bc)}}, bc), T) + + # this doesn't escape for Array but for JLArray return Grid(Array{T}(undef, size(bc)...)) end diff --git a/test/grids.jl b/test/grids.jl index 1eb1a86f0..3336ed9ee 100644 --- a/test/grids.jl +++ b/test/grids.jl @@ -1,3 +1,6 @@ +using JLArrays +using Adapt + @testset "Grid indexing" begin for NF in (Float32, Float64) @@ -323,8 +326,6 @@ end end end -using JLArrays - @testset "AbstractGridArray: GPU (JLArrays)" begin NF = Float32 @testset for Grid in ( From cfae03427b6589373896af7b9f4f766c747bc74c Mon Sep 17 00:00:00 2001 From: Milan Date: Wed, 17 Apr 2024 19:36:21 -0400 Subject: [PATCH 16/35] full_grid_type vs full_array_type plus tests --- src/RingGrids/full_grids.jl | 3 ++- src/output/output.jl | 4 ++-- test/grids.jl | 42 +++++++++++++++++++++++++++++++++++++ 3 files changed, 46 insertions(+), 3 deletions(-) diff --git a/src/RingGrids/full_grids.jl b/src/RingGrids/full_grids.jl index c63d9ccb0..685079ec5 100644 --- a/src/RingGrids/full_grids.jl +++ b/src/RingGrids/full_grids.jl @@ -7,7 +7,8 @@ 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.""" const AbstractFullGrid{T} = AbstractFullGridArray{T, 1, Vector{T}} -full_grid_type(Grid::Type{<:AbstractFullGridArray}) = nonparametric_type(Grid) +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) diff --git a/src/output/output.jl b/src/output/output.jl index 4dc815619..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" @@ -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 3336ed9ee..0be1357ef 100644 --- a/test/grids.jl +++ b/test/grids.jl @@ -1,6 +1,48 @@ 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) From 38d40627a9526fad940f4311be6545f1e46c89c5 Mon Sep 17 00:00:00 2001 From: Milan Date: Wed, 17 Apr 2024 19:46:08 -0400 Subject: [PATCH 17/35] use octahealpix weights for octahealpix not healpix --- src/RingGrids/grids/full_octahealpix.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/RingGrids/grids/full_octahealpix.jl b/src/RingGrids/grids/full_octahealpix.jl index 52296056f..b38a52792 100644 --- a/src/RingGrids/grids/full_octahealpix.jl +++ b/src/RingGrids/grids/full_octahealpix.jl @@ -25,7 +25,7 @@ get_colat(::Type{<:FullOctaHEALPixArray}, nlat_half::Integer) = get_colat(OctaHE get_lon(::Type{<:FullOctaHEALPixArray}, nlat_half::Integer) = get_lon(FullGaussianArray, nlat_half) # QUADRATURE -get_quadrature_weights(::Type{<:FullOctaHEALPixArray}, nlat_half::Integer) = healpix_weights(nlat_half) +get_quadrature_weights(::Type{<:FullOctaHEALPixArray}, nlat_half::Integer) = octahealpix_weights(nlat_half) """ G = FullOctaHEALPixGrid{T} From 891363f82fac74149b83fb57a349b87c2038bff7 Mon Sep 17 00:00:00 2001 From: Milan Date: Wed, 17 Apr 2024 19:55:57 -0400 Subject: [PATCH 18/35] n->nlat_half typo --- src/RingGrids/grids/octahedral_clenshaw.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/RingGrids/grids/octahedral_clenshaw.jl b/src/RingGrids/grids/octahedral_clenshaw.jl index 593c9a88c..c676ca163 100644 --- a/src/RingGrids/grids/octahedral_clenshaw.jl +++ b/src/RingGrids/grids/octahedral_clenshaw.jl @@ -42,7 +42,7 @@ matrix_size(grid::Grid) where {Grid<:OctahedralClenshawGrid} = matrix_size(Grid, 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 + 4n)÷2 + N = (o + 4nlat_half)÷2 return (N, N) end From f819cc1803f091bd6c4c4f236f739cfc7a498ebf Mon Sep 17 00:00:00 2001 From: Milan Date: Wed, 17 Apr 2024 23:56:17 -0400 Subject: [PATCH 19/35] docstrings for general and full grids --- src/RingGrids/full_grids.jl | 22 +++-- src/RingGrids/general.jl | 111 +++++++++++++++++------- src/RingGrids/grids/full_clenshaw.jl | 25 +++--- src/RingGrids/grids/full_gaussian.jl | 24 ++++- src/RingGrids/grids/full_healpix.jl | 25 +++--- src/RingGrids/grids/full_octahealpix.jl | 24 ++--- src/RingGrids/reduced_grids.jl | 12 +++ 7 files changed, 167 insertions(+), 76 deletions(-) diff --git a/src/RingGrids/full_grids.jl b/src/RingGrids/full_grids.jl index 685079ec5..2b9404305 100644 --- a/src/RingGrids/full_grids.jl +++ b/src/RingGrids/full_grids.jl @@ -1,11 +1,11 @@ +"""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 -""" -abstract type AbstractFullGrid{T} <: AbstractGrid{T} 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.""" +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) @@ -24,18 +24,22 @@ Base.Array(grid::AbstractFullGridArray) = Array(reshape(grid.data, :, get_nlat(g 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 = 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 + 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 +# precompute ring indices for full grids function each_index_in_ring!( rings::Vector{<:UnitRange{<:Integer}}, Grid::Type{<:AbstractFullGridArray}, diff --git a/src/RingGrids/general.jl b/src/RingGrids/general.jl index f57c03ab1..34e26ee7f 100644 --- a/src/RingGrids/general.jl +++ b/src/RingGrids/general.jl @@ -1,33 +1,69 @@ +"""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 -# Horizontal 2D grids with N=1 +"""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)) + +"""$(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) -Base.size(G::AbstractGridArray) = size(G.data) -Base.sizeof(G::AbstractGridArray) = sizeof(G.data) +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) -# get total number of grid points vs horizontal (2D) only +"""$(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) -# size of the matrix of the horizontal grid if representable as such (not all grids) +"""$(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 @@ -38,7 +74,7 @@ matrix_size(grid::Grid) where {Grid<:AbstractGridArray} = matrix_size(Grid, grid col::Colon, k::Integer..., ) where {GridArray<:AbstractGridArray} - GridArray_ = nonparametric_type(GridArray) + GridArray_ = nonparametric_type(GridArray) # obtain parameters from G.data return GridArray_(getindex(G.data, col, k...), G.nlat_half, G.rings) end @@ -50,6 +86,8 @@ end setindex!(G.data, x, ij, k) ## 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 @@ -71,15 +109,15 @@ end # if no rings are provided, calculate them function (::Type{Grid})(data::AbstractArray, nlat_half::Integer) where {Grid<:AbstractGridArray} - GridArray_ = nonparametric_type(Grid) - rings = eachring(Grid, nlat_half) + 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) - nlat_half = get_nlat_half(Grid, npoints2D) + 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 @@ -177,6 +215,9 @@ function (::Type{Grid})( 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) @@ -189,6 +230,10 @@ function get_latdlonds(Grid::Type{<:AbstractGridArray}, nlat_half::Integer) 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 @@ -196,11 +241,23 @@ function get_latlons(Grid::Type{<:AbstractGridArray}, nlat_half::Integer) 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) @@ -212,10 +269,6 @@ function get_latd(Grid::Type{<:AbstractGridArray}, nlat_half::Integer) return get_lat(Grid, nlat_half) * (360/2π) end -# only defined for full grids, empty vector as fallback -get_lon(::Type{<:AbstractGridArray}, nlat_half::Integer) = Float64[] -get_lond(::Type{<:AbstractGridArray}, nlat_half::Integer) = Float64[] - """ $(TYPEDSIGNATURES) Returns a vector `nlons` for the number of longitude points per latitude ring, north to south. @@ -250,40 +303,38 @@ and then each grid point per ring. To be used like @inline eachring(grid::AbstractGridArray) = grid.rings function eachring(Grid::Type{<:AbstractGridArray}, nlat_half::Integer) - rings = Vector{UnitRange{Int}}(undef, get_nlat(Grid, nlat_half)) - each_index_in_ring!(rings, Grid, nlat_half) + 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 -in `grids` are of same size.""" +"""$(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 +"""$(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`.""" +"""$(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`.""" +""" $(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) @@ -291,10 +342,8 @@ function eachgridpoint(grid1::Grid, grids::Grid...) where {Grid<:AbstractGridArr return eachgridpoint(grid1) end -""" -$(TYPEDSIGNATURES) -Obtain ring index j from gridpoint ij and Vector{UnitRange} describing rind indices -as obtained from eachring(::Grid)""" +"""$(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 diff --git a/src/RingGrids/grids/full_clenshaw.jl b/src/RingGrids/grids/full_clenshaw.jl index 6d9fce596..2e10643f4 100644 --- a/src/RingGrids/grids/full_clenshaw.jl +++ b/src/RingGrids/grids/full_clenshaw.jl @@ -1,3 +1,13 @@ +"""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 @@ -14,6 +24,9 @@ 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 @@ -25,14 +38,4 @@ get_colat(::Type{<:FullClenshawArray}, nlat_half::Integer) = [j/(2nlat_half)*π 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) - - -""" - 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.""" -FullClenshawGrid \ No newline at end of file +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 index 02448c5c0..044aed48c 100644 --- a/src/RingGrids/grids/full_gaussian.jl +++ b/src/RingGrids/grids/full_gaussian.jl @@ -1,6 +1,19 @@ +"""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::ArrayType # data array, ring by ring, north to south - nlat_half::Int # number of latitudes on one hemisphere + "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} = @@ -11,11 +24,14 @@ end # TYPES const FullGaussianGrid{T} = FullGaussianArray{T, 1, Vector{T}} -nonparametric_type(::Type{<:FullGaussianArray}) = FullGaussianArray +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 +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 diff --git a/src/RingGrids/grids/full_healpix.jl b/src/RingGrids/grids/full_healpix.jl index 045f1ee43..abb7e59bc 100644 --- a/src/RingGrids/grids/full_healpix.jl +++ b/src/RingGrids/grids/full_healpix.jl @@ -1,3 +1,13 @@ +"""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 @@ -14,6 +24,9 @@ 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) @@ -25,14 +38,4 @@ get_colat(::Type{<:FullHEALPixArray}, nlat_half::Integer) = get_colat(HEALPixGri get_lon(::Type{<:FullHEALPixArray}, nlat_half::Integer) = get_lon(FullGaussianArray, nlat_half) # QUADRATURE -get_quadrature_weights(::Type{<:FullHEALPixArray}, nlat_half::Integer) = healpix_weights(nlat_half) - - -""" - 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.""" -FullHEALPixGrid \ No newline at end of file +get_quadrature_weights(::Type{<:FullHEALPixArray}, nlat_half::Integer) = healpix_weights(nlat_half) \ No newline at end of file diff --git a/src/RingGrids/grids/full_octahealpix.jl b/src/RingGrids/grids/full_octahealpix.jl index b38a52792..8c19e3f20 100644 --- a/src/RingGrids/grids/full_octahealpix.jl +++ b/src/RingGrids/grids/full_octahealpix.jl @@ -1,3 +1,13 @@ +"""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 @@ -14,6 +24,9 @@ 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) @@ -25,13 +38,4 @@ get_colat(::Type{<:FullOctaHEALPixArray}, nlat_half::Integer) = get_colat(OctaHE get_lon(::Type{<:FullOctaHEALPixArray}, nlat_half::Integer) = get_lon(FullGaussianArray, nlat_half) # QUADRATURE -get_quadrature_weights(::Type{<:FullOctaHEALPixArray}, nlat_half::Integer) = octahealpix_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.""" -FullOctaHEALPixGrid \ No newline at end of file +get_quadrature_weights(::Type{<:FullOctaHEALPixArray}, nlat_half::Integer) = octahealpix_weights(nlat_half) \ No newline at end of file diff --git a/src/RingGrids/reduced_grids.jl b/src/RingGrids/reduced_grids.jl index c41571b6e..37985935a 100644 --- a/src/RingGrids/reduced_grids.jl +++ b/src/RingGrids/reduced_grids.jl @@ -1,6 +1,18 @@ +"""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) From 5f5db143b15d999180101699e1034ffd0d9fa390 Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 18 Apr 2024 19:45:24 +0200 Subject: [PATCH 20/35] GPU broadcast fix? --- src/RingGrids/general.jl | 27 ++++++++++++--------------- 1 file changed, 12 insertions(+), 15 deletions(-) diff --git a/src/RingGrids/general.jl b/src/RingGrids/general.jl index 34e26ee7f..e014d0b34 100644 --- a/src/RingGrids/general.jl +++ b/src/RingGrids/general.jl @@ -369,13 +369,6 @@ Base.BroadcastStyle(::Type{Grid}) where {Grid<:AbstractGridArray{T, N, ArrayType # 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} - # this escapes for Array and JLArray - # if isstructurepreserving(bc) || fzeropreserving(bc) - # return Grid(Array{T}(undef, size(bc)...)) - # end - # return similar(convert(Broadcasted{DefaultArrayStyle{ndims(bc)}}, bc), T) - - # this doesn't escape for Array but for JLArray return Grid(Array{T}(undef, size(bc)...)) end @@ -383,18 +376,22 @@ end AbstractGridArrayStyle{N, Grid}(::Val{M}) where {N, Grid, M} = AbstractGridArrayStyle{N, Grid}() ## GPU -function Broadcast.BroadcastStyle( - ::Type{Grid}, -) where {Grid <: AbstractGridArray{T, N, ArrayType}} where {T, N, ArrayType <: GPUArrays.AbstractGPUArray} - return Broadcast.BroadcastStyle(ArrayType) -end +struct AbstractGPUGridArrayStyle{N, ArrayType, Grid} <: GPUArrays.AbstractGPUArrayStyle{N} end -function GPUArrays.backend( - ::Type{Grid}, -) where {Grid <: AbstractGridArray{T, N, ArrayType}} where {T, N, ArrayType <: GPUArrays.AbstractGPUArray} +Base.BroadcastStyle(::Type{Grid}) where {Grid<:AbstractGridArray{T, N, ArrayType}} where {T, N, ArrayType <: GPUArrays.AbstractGPUArray} = + AbstractGPUGridArrayStyle{N, ArrayType, nonparametric_type(Grid)}() + +# ::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} + return Grid(T.(ArrayType(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) From 7b5ad57ebe9decf769a4b2e29ee77657631f12e4 Mon Sep 17 00:00:00 2001 From: Milan Date: Thu, 18 Apr 2024 14:19:06 -0400 Subject: [PATCH 21/35] Reduced grid documentation --- src/RingGrids/grids/healpix.jl | 54 +++++++++++++++++----- src/RingGrids/grids/new_grids.jl | 42 ----------------- src/RingGrids/grids/octahealpix.jl | 39 ++++++++++++---- src/RingGrids/grids/octahedral_clenshaw.jl | 50 +++++++++++--------- src/RingGrids/grids/octahedral_gaussian.jl | 46 ++++++++++++------ 5 files changed, 134 insertions(+), 97 deletions(-) delete mode 100644 src/RingGrids/grids/new_grids.jl diff --git a/src/RingGrids/grids/healpix.jl b/src/RingGrids/grids/healpix.jl index 6bd47058c..c8f871495 100644 --- a/src/RingGrids/grids/healpix.jl +++ b/src/RingGrids/grids/healpix.jl @@ -1,3 +1,23 @@ +"""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 @@ -15,17 +35,33 @@ 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 -nside_healpix(nlat_half::Integer) = nlat_half÷2 +"""$(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) @@ -39,6 +75,7 @@ function get_colat(::Type{<:HEALPixArray}, nlat_half::Integer) 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 @@ -50,6 +87,7 @@ function get_lon_per_ring(Grid::Type{<:HEALPixArray}, nlat_half::Integer, j::Int 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] @@ -87,6 +125,9 @@ function each_index_in_ring!( rings::Vector{<:UnitRange{<:Integer}}, 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 @@ -114,13 +155,4 @@ function each_index_in_ring!( rings::Vector{<:UnitRange{<:Integer}}, index_end += 4j_mirrored # 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.""" -HEALPixGrid \ No newline at end of file +end \ No newline at end of file diff --git a/src/RingGrids/grids/new_grids.jl b/src/RingGrids/grids/new_grids.jl deleted file mode 100644 index f2cfc36dc..000000000 --- a/src/RingGrids/grids/new_grids.jl +++ /dev/null @@ -1,42 +0,0 @@ -""" - 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 -""" \ No newline at end of file diff --git a/src/RingGrids/grids/octahealpix.jl b/src/RingGrids/grids/octahealpix.jl index 6287cd736..bcdf5b1eb 100644 --- a/src/RingGrids/grids/octahealpix.jl +++ b/src/RingGrids/grids/octahealpix.jl @@ -1,3 +1,24 @@ +"""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 @@ -15,10 +36,15 @@ nonparametric_type(::Type{<:OctaHEALPixArray}) = OctaHEALPixArray horizontal_grid_type(::Type{<:OctaHEALPixArray}) = OctaHEALPixGrid full_array_type(::Type{<:OctaHEALPixArray}) = FullOctaHEALPixArray +"""An `OctaHEALPixArray` but constrained to `N=1` dimensions (horizontal only) and data is a `Vector{T}`.""" +OctaHEALPixGrid + ## 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)) + +# 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." @@ -33,6 +59,7 @@ 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 @@ -41,6 +68,8 @@ 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 @@ -167,12 +196,4 @@ function Matrix!( MGs::Tuple{AbstractMatrix{T}, OctaHEALPixGrid}...; ntuples == 1 && return M return Tuple(Mi for (Mi, Gi) in MGs) -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.""" -OctaHEALPixGrid \ No newline at end of file +end \ No newline at end of file diff --git a/src/RingGrids/grids/octahedral_clenshaw.jl b/src/RingGrids/grids/octahedral_clenshaw.jl index c676ca163..1314d6c54 100644 --- a/src/RingGrids/grids/octahedral_clenshaw.jl +++ b/src/RingGrids/grids/octahedral_clenshaw.jl @@ -1,3 +1,21 @@ +"""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 @@ -15,6 +33,9 @@ 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 @@ -46,7 +67,8 @@ function matrix_size(::Type{OctahedralClenshawGrid}, nlat_half::Integer) return (N, N) end -Base.Matrix(G::OctahedralClenshawGrid{T}; kwargs...) where T = Matrix!(zeros(T, matrix_size(G)...), G; kwargs...) +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) @@ -56,7 +78,8 @@ function get_lon_per_ring(Grid::Type{<:OctahedralClenshawArray}, nlat_half::Inte end ## QUADRATURE -get_quadrature_weights(::Type{<:OctahedralClenshawArray}, nlat_half::Integer) = clenshaw_curtis_weights(nlat_half) +get_quadrature_weights(::Type{<:OctahedralClenshawArray}, nlat_half::Integer) = + clenshaw_curtis_weights(nlat_half) ## INDEXING function each_index_in_ring(Grid::Type{<:OctahedralClenshawArray}, @@ -107,12 +130,7 @@ end ## CONVERSION """ - Matrix!(M::AbstractMatrix, - G::OctahedralClenshawGrid; - quadrant_rotation=(0, 1, 2, 3), - matrix_quadrant=((2, 2), (1, 2), (1, 1), (2, 1)), - ) - +$(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 @@ -122,8 +140,7 @@ 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...) - +$(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.""" @@ -190,15 +207,4 @@ function Matrix!( MGs::Tuple{AbstractMatrix{T}, OctahedralClenshawGrid}...; ntuples == 1 && return M return Tuple(Mi for (Mi, Gi) in MGs) -end - -""" - 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.""" \ No newline at end of file +end \ No newline at end of file diff --git a/src/RingGrids/grids/octahedral_gaussian.jl b/src/RingGrids/grids/octahedral_gaussian.jl index 7a33b36d8..b6246e2f6 100644 --- a/src/RingGrids/grids/octahedral_gaussian.jl +++ b/src/RingGrids/grids/octahedral_gaussian.jl @@ -1,3 +1,20 @@ +"""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 @@ -15,9 +32,18 @@ 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) @@ -52,6 +78,9 @@ end 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 @@ -75,6 +104,9 @@ function each_index_in_ring(Grid::Type{<:OctahedralGaussianArray}, 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 @@ -96,16 +128,4 @@ function each_index_in_ring!( rings::Vector{<:UnitRange{<:Integer}}, index_end += o + m*j_mirrored # 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. -# """ -# # OctahedralGaussianGrid \ No newline at end of file +end \ No newline at end of file From d63431620eb140ba7c50976937f1a0a488015ee1 Mon Sep 17 00:00:00 2001 From: Milan Date: Fri, 19 Apr 2024 18:02:16 -0400 Subject: [PATCH 22/35] Add JLArrays extension --- Project.toml | 3 +++ ext/SpeedyWeatherJLArraysExt.jl | 8 ++++++++ src/RingGrids/RingGrids.jl | 2 +- src/RingGrids/general.jl | 7 ++++++- 4 files changed, 18 insertions(+), 2 deletions(-) create mode 100644 ext/SpeedyWeatherJLArraysExt.jl diff --git a/Project.toml b/Project.toml index c952774f1..046cb0a7b 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" @@ -49,6 +51,7 @@ 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" diff --git a/ext/SpeedyWeatherJLArraysExt.jl b/ext/SpeedyWeatherJLArraysExt.jl new file mode 100644 index 000000000..6d1346251 --- /dev/null +++ b/ext/SpeedyWeatherJLArraysExt.jl @@ -0,0 +1,8 @@ +module SpeedyWeatherMakieExt + +using SpeedyWeather, JLArrays + +# for RingGrids, every Array needs this method to strip away the parameters +SpeedyWeather.RingGrids.nonparametric_type(::Type{JLArray}) = JLArray + +end # module \ No newline at end of file diff --git a/src/RingGrids/RingGrids.jl b/src/RingGrids/RingGrids.jl index 7a6ffcbc7..7fb68f35d 100644 --- a/src/RingGrids/RingGrids.jl +++ b/src/RingGrids/RingGrids.jl @@ -11,8 +11,8 @@ import LinearAlgebra # GPU import Adapt -import KernelAbstractions import GPUArrays +import CUDA # GRIDS export AbstractGridArray, diff --git a/src/RingGrids/general.jl b/src/RingGrids/general.jl index e014d0b34..08647799c 100644 --- a/src/RingGrids/general.jl +++ b/src/RingGrids/general.jl @@ -19,6 +19,10 @@ const AbstractGrid{T} = AbstractGridArray{T, 1, Vector{T}} (*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{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.""" @@ -193,7 +197,8 @@ function (::Type{Grid})( nlat_half::Integer, k::Integer..., ) where {Grid<:AbstractGridArray{T, N, ArrayType}} where {T, N, ArrayType} - return Grid(ArrayType{T, N}(undef, get_npoints2D(Grid, nlat_half), k...), nlat_half) + 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 From 9ea9440d3b9033b0bca3c9b6e3d9a1fff2dd6727 Mon Sep 17 00:00:00 2001 From: Milan Date: Fri, 19 Apr 2024 18:06:37 -0400 Subject: [PATCH 23/35] I can't close "" --- Project.toml | 2 +- src/RingGrids/RingGrids.jl | 7 +++---- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/Project.toml b/Project.toml index 046cb0a7b..2ad23b26f 100644 --- a/Project.toml +++ b/Project.toml @@ -35,7 +35,7 @@ JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb" [extensions] SpeedyWeatherMakieExt = "Makie" -SpeedyWeatherJLArraysExt = "JLArrays +SpeedyWeatherJLArraysExt = "JLArrays" [compat] AbstractFFTs = "1" diff --git a/src/RingGrids/RingGrids.jl b/src/RingGrids/RingGrids.jl index 7fb68f35d..a5d46952a 100644 --- a/src/RingGrids/RingGrids.jl +++ b/src/RingGrids/RingGrids.jl @@ -14,17 +14,16 @@ import Adapt import GPUArrays import CUDA -# GRIDS +# 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, From 5b087e919584dd37cd77d3bd29eeb437d14210bf Mon Sep 17 00:00:00 2001 From: Milan Date: Fri, 19 Apr 2024 18:27:47 -0400 Subject: [PATCH 24/35] CUDA.CuArray for nonparametric_type --- src/RingGrids/general.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/RingGrids/general.jl b/src/RingGrids/general.jl index 08647799c..98b298f48 100644 --- a/src/RingGrids/general.jl +++ b/src/RingGrids/general.jl @@ -21,7 +21,7 @@ nonparametric_type(grid::AbstractGridArray) = nonparametric_type(typeof(grid)) # also needed for other array types nonparametric_type(::Type{Array}) = Array -nonparametric_type(::Type{CuArray}) = CuArray +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 From fd6d9d8fdf7065577be505ee729d1b3a0dc24093 Mon Sep 17 00:00:00 2001 From: Milan Date: Sat, 20 Apr 2024 12:46:12 -0400 Subject: [PATCH 25/35] use nonparametric_type(ArrayType) to avoid conversion --- src/RingGrids/general.jl | 22 ++++++++++++++++------ 1 file changed, 16 insertions(+), 6 deletions(-) diff --git a/src/RingGrids/general.jl b/src/RingGrids/general.jl index 98b298f48..4d25caa1a 100644 --- a/src/RingGrids/general.jl +++ b/src/RingGrids/general.jl @@ -383,18 +383,28 @@ AbstractGridArrayStyle{N, Grid}(::Val{M}) where {N, Grid, M} = AbstractGridArray ## GPU struct AbstractGPUGridArrayStyle{N, ArrayType, Grid} <: GPUArrays.AbstractGPUArrayStyle{N} end -Base.BroadcastStyle(::Type{Grid}) where {Grid<:AbstractGridArray{T, N, ArrayType}} where {T, N, ArrayType <: GPUArrays.AbstractGPUArray} = - AbstractGPUGridArrayStyle{N, ArrayType, nonparametric_type(Grid)}() +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}() +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} +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} - return Grid(T.(ArrayType(undef, size(bc)...))) +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} From 30ba3cec3938419bf5d1ed22a66b186f2c8e687d Mon Sep 17 00:00:00 2001 From: Milan Date: Sat, 20 Apr 2024 12:55:55 -0400 Subject: [PATCH 26/35] JLArrays extension was incorrectly named --- ext/SpeedyWeatherJLArraysExt.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ext/SpeedyWeatherJLArraysExt.jl b/ext/SpeedyWeatherJLArraysExt.jl index 6d1346251..c80745ec2 100644 --- a/ext/SpeedyWeatherJLArraysExt.jl +++ b/ext/SpeedyWeatherJLArraysExt.jl @@ -1,4 +1,4 @@ -module SpeedyWeatherMakieExt +module SpeedyWeatherJLArraysExt using SpeedyWeather, JLArrays From 028230862a81c39542966cc86ec7fe234fb6fb54 Mon Sep 17 00:00:00 2001 From: Milan Date: Mon, 22 Apr 2024 15:37:32 -0400 Subject: [PATCH 27/35] quadrature weights of size nlat --- src/RingGrids/general.jl | 6 +-- src/RingGrids/grids/full_healpix.jl | 2 +- src/RingGrids/grids/full_octahealpix.jl | 3 +- src/RingGrids/interpolation.jl | 3 +- src/RingGrids/quadrature_weights.jl | 51 ++++++++++++++++------ src/SpeedyTransforms/spectral_transform.jl | 3 +- 6 files changed, 46 insertions(+), 22 deletions(-) diff --git a/src/RingGrids/general.jl b/src/RingGrids/general.jl index 4d25caa1a..653b19730 100644 --- a/src/RingGrids/general.jl +++ b/src/RingGrids/general.jl @@ -277,9 +277,9 @@ 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{<:AbstractGridArray}, nlat_half::Integer; both_hemispheres::Bool=false) +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 diff --git a/src/RingGrids/grids/full_healpix.jl b/src/RingGrids/grids/full_healpix.jl index abb7e59bc..1c483767b 100644 --- a/src/RingGrids/grids/full_healpix.jl +++ b/src/RingGrids/grids/full_healpix.jl @@ -38,4 +38,4 @@ get_colat(::Type{<:FullHEALPixArray}, nlat_half::Integer) = get_colat(HEALPixGri get_lon(::Type{<:FullHEALPixArray}, nlat_half::Integer) = get_lon(FullGaussianArray, nlat_half) # QUADRATURE -get_quadrature_weights(::Type{<:FullHEALPixArray}, nlat_half::Integer) = healpix_weights(nlat_half) \ No newline at end of file +get_quadrature_weights(::Type{<:FullHEALPixArray}, nlat_half::Integer) = equal_area_weights(FullHEALPixArray, nlat_half) \ No newline at end of file diff --git a/src/RingGrids/grids/full_octahealpix.jl b/src/RingGrids/grids/full_octahealpix.jl index 8c19e3f20..ccd5838c7 100644 --- a/src/RingGrids/grids/full_octahealpix.jl +++ b/src/RingGrids/grids/full_octahealpix.jl @@ -38,4 +38,5 @@ get_colat(::Type{<:FullOctaHEALPixArray}, nlat_half::Integer) = get_colat(OctaHE get_lon(::Type{<:FullOctaHEALPixArray}, nlat_half::Integer) = get_lon(FullGaussianArray, nlat_half) # QUADRATURE -get_quadrature_weights(::Type{<:FullOctaHEALPixArray}, nlat_half::Integer) = octahealpix_weights(nlat_half) \ No newline at end of file +get_quadrature_weights(::Type{<:FullOctaHEALPixArray}, nlat_half::Integer) = + equal_area_weights(FullOctaHEALPixArray, nlat_half) \ No newline at end of file diff --git a/src/RingGrids/interpolation.jl b/src/RingGrids/interpolation.jl index 7ebed411e..06480b492 100644 --- a/src/RingGrids/interpolation.jl +++ b/src/RingGrids/interpolation.jl @@ -46,8 +46,7 @@ function GridGeometry( Grid::Type{<:AbstractGrid}, # which grid to calculate th # 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) + 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) diff --git a/src/RingGrids/quadrature_weights.jl b/src/RingGrids/quadrature_weights.jl index 983728f2b..dc997557c 100644 --- a/src/RingGrids/quadrature_weights.jl +++ b/src/RingGrids/quadrature_weights.jl @@ -1,28 +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 = 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 +"""$(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) - weights = zeros(nlat_half) - for j in 1:nlat_half + 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/get_npoints2D(Grid, nlat_half) + weights[j] = 2nlon/npoints2D end return weights end -healpix_weights(nlat_half::Integer) = equal_area_weights(HEALPixArray, nlat_half) -octahealpix_weights(nlat_half::Integer) = equal_area_weights(OctaHEALPixArray, nlat_half) - # SOLID ANGLES ΔΩ = sinθ Δθ Δϕ -get_solid_angles(Grid::Type{<:AbstractGridArray}, 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{HEALPixArray, OctaHEALPixArray}}, nlat_half::Integer) = - 4π/get_npoints(Grid, nlat_half)*ones(nlat_half) + 4π/get_npoints2D(Grid, nlat_half)*ones(get_nlat(Grid, nlat_half)) \ No newline at end of file diff --git a/src/SpeedyTransforms/spectral_transform.jl b/src/SpeedyTransforms/spectral_transform.jl index 69a27bc87..28efefff8 100644 --- a/src/SpeedyTransforms/spectral_transform.jl +++ b/src/SpeedyTransforms/spectral_transform.jl @@ -39,7 +39,8 @@ struct SpectralTransform{NF<:AbstractFloat} Λs::Vector{LowerTriangularMatrix{NF}} # Legendre polynomials for all latitudes (all precomputed) # SOLID ANGLES ΔΩ FOR QUADRATURE - # (integration for the Legendre polynomials, extra normalisation of π/nlat included) + # (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 From dc8fa568fbe429b2f949063a0a32c459249dfc28 Mon Sep 17 00:00:00 2001 From: Milan Date: Mon, 22 Apr 2024 17:35:38 -0400 Subject: [PATCH 28/35] AbstractGridArray methods to avoid scalar indexing, @propagate_inbounds corrected --- src/RingGrids/general.jl | 16 +++++++----- test/grids.jl | 56 +++++++++++++++++++++++++--------------- 2 files changed, 44 insertions(+), 28 deletions(-) diff --git a/src/RingGrids/general.jl b/src/RingGrids/general.jl index 653b19730..ee4600d3f 100644 --- a/src/RingGrids/general.jl +++ b/src/RingGrids/general.jl @@ -71,7 +71,8 @@ get_npoints2D(grid::Grid) where {Grid<:AbstractGridArray} = get_npoints2D(Grid, matrix_size(grid::Grid) where {Grid<:AbstractGridArray} = matrix_size(Grid, grid.nlat_half) ## INDEXING -@inline Base.getindex(G::AbstractGridArray, ijk...) = getindex(G.data, ijk...) +# simply propagate all indices forward +Base.@propagate_inbounds Base.getindex(G::AbstractGridArray, ijk...) = getindex(G.data, ijk...) @inline function Base.getindex( G::GridArray, @@ -82,12 +83,9 @@ matrix_size(grid::Grid) where {Grid<:AbstractGridArray} = matrix_size(Grid, grid return GridArray_(getindex(G.data, col, k...), G.nlat_half, G.rings) end -@inline Base.setindex!(G::AbstractGridArray, x, ijk::Integer...) = - setindex!(G.data, x, ijk...) -@inline Base.setindex!(G::AbstractGridArray, x::AbstractVector, ij::AbstractRange, k::Integer...) = - setindex!(G.data, x, ij, k...) -@inline Base.setindex!(G::AbstractGridArray, x::AbstractVector, ij::Integer, k::AbstractRange) = - setindex!(G.data, x, ij, k) +# 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 @@ -321,6 +319,10 @@ function eachring(grid1::Grid, grids::Grid...) where {Grid<:AbstractGridArray} return eachring(grid1) end +# equality +Base.:(==)(G1::AbstractGridArray, G2::AbstractGridArray) = grids_match(G1, G2) && G1.data == G2.data +Base.all(G::AbstractGridArray) = all(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.""" diff --git a/test/grids.jl b/test/grids.jl index 0be1357ef..568c188e7 100644 --- a/test/grids.jl +++ b/test/grids.jl @@ -368,6 +368,8 @@ end end end +RingGrids.nonparametric_type(::Type{<:JLArray}) = JLArray + @testset "AbstractGridArray: GPU (JLArrays)" begin NF = Float32 @testset for Grid in ( @@ -380,32 +382,44 @@ end FullHEALPixArray, FullOctaHEALPixArray, ) - G_cpu = randn(Grid{NF}, 2, 3, 4) + 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 - # @test all(G .== G2) + @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 - # @test G[:, 1] isa Grid{NF, 1, JLArray{NF}} - # for k in eachgrid(G) - # for (j, ring) in enumerate(eachring(G)) - # @test Array(L[lm,:]) == L_cpu[lm,:] - # end - # end - - # # setindex! - # A_test = JLArray(rand(NF,size(L_cpu,3))) - # L[1,:] = A_test - # @test L[1,:] == A_test - - # # fill - # fill!(L, 2) - # for lm in SpeedyWeather.eachharmonic(L2) - # @test all(L[lm, [Colon() for i=1:length(idims)]...] .== 2) - # end + @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 + + + # fill + fill!(G, 2) + @test all(G .== 2) end end \ No newline at end of file From 61dc760c3fb247a38b7f92c4ad0d21222919c1cb Mon Sep 17 00:00:00 2001 From: Milan Date: Mon, 22 Apr 2024 17:42:16 -0400 Subject: [PATCH 29/35] more setindex tests for AbstractGridArray --- test/grids.jl | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/test/grids.jl b/test/grids.jl index 568c188e7..3a300bce2 100644 --- a/test/grids.jl +++ b/test/grids.jl @@ -417,6 +417,15 @@ RingGrids.nonparametric_type(::Type{<:JLArray}) = JLArray 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) From 1ca3a906244b426cdea5a63835bcc776c5bbb267 Mon Sep 17 00:00:00 2001 From: Milan Date: Mon, 22 Apr 2024 18:01:57 -0400 Subject: [PATCH 30/35] SpectralTransform uses nonparametric type of AbstractGridArray --- src/SpeedyTransforms/spectral_transform.jl | 62 +++++++++++----------- src/dynamics/spectral_grid.jl | 2 +- 2 files changed, 33 insertions(+), 31 deletions(-) diff --git a/src/SpeedyTransforms/spectral_transform.jl b/src/SpeedyTransforms/spectral_transform.jl index cb5fe3c12..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 @@ -71,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) @@ -117,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 @@ -131,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))) @@ -153,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 @@ -189,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] @@ -360,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 @@ -420,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 @@ -461,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 @@ -504,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) @@ -541,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) @@ -559,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/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) From b9bd961ba8c09dc1f99c7ddf26474b1390ed3656 Mon Sep 17 00:00:00 2001 From: Milan Date: Mon, 22 Apr 2024 18:39:06 -0400 Subject: [PATCH 31/35] corrected quadrature weights for full healpix grids --- src/RingGrids/grids/full_healpix.jl | 4 ++-- src/RingGrids/grids/full_octahealpix.jl | 4 ++-- test/grids.jl | 2 ++ 3 files changed, 6 insertions(+), 4 deletions(-) diff --git a/src/RingGrids/grids/full_healpix.jl b/src/RingGrids/grids/full_healpix.jl index 1c483767b..8a6aefed4 100644 --- a/src/RingGrids/grids/full_healpix.jl +++ b/src/RingGrids/grids/full_healpix.jl @@ -37,5 +37,5 @@ get_nlon(::Type{<:FullHEALPixArray}, nlat_half::Integer) = 4nlat_half 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 -get_quadrature_weights(::Type{<:FullHEALPixArray}, nlat_half::Integer) = equal_area_weights(FullHEALPixArray, nlat_half) \ No newline at end of file +# 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 index ccd5838c7..1e604bfcf 100644 --- a/src/RingGrids/grids/full_octahealpix.jl +++ b/src/RingGrids/grids/full_octahealpix.jl @@ -37,6 +37,6 @@ get_nlon(::Type{<:FullOctaHEALPixArray}, nlat_half::Integer) = 4nlat_half 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 +# QUADRATURE (use weights from reduced grids though!) get_quadrature_weights(::Type{<:FullOctaHEALPixArray}, nlat_half::Integer) = - equal_area_weights(FullOctaHEALPixArray, nlat_half) \ No newline at end of file + equal_area_weights(OctaHEALPixArray, nlat_half) \ No newline at end of file diff --git a/test/grids.jl b/test/grids.jl index 3a300bce2..8249b6529 100644 --- a/test/grids.jl +++ b/test/grids.jl @@ -279,7 +279,9 @@ end 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 From 77a759241733f7b7896ac58526b2b82b308d3844 Mon Sep 17 00:00:00 2001 From: Milan Date: Tue, 23 Apr 2024 10:02:39 -0400 Subject: [PATCH 32/35] any(::AbstractGridArray) added --- src/RingGrids/general.jl | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/src/RingGrids/general.jl b/src/RingGrids/general.jl index ee4600d3f..b26515382 100644 --- a/src/RingGrids/general.jl +++ b/src/RingGrids/general.jl @@ -302,9 +302,15 @@ and then each grid point per ring. To be used like rings = eachring(grid) for ring in rings for ij in ring - grid[ij]""" + 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 @@ -319,9 +325,10 @@ function eachring(grid1::Grid, grids::Grid...) where {Grid<:AbstractGridArray} return eachring(grid1) end -# equality +# 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 From 7f4bf3812b1409dc36f0e254283dd71cbecc5f96 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Milan=20Kl=C3=B6wer?= Date: Thu, 25 Apr 2024 11:56:40 -0400 Subject: [PATCH 33/35] nonparametric_type in LTA for JLArrays moved to extension --- ext/SpeedyWeatherJLArraysExt.jl | 4 +++- test/lower_triangular_matrix.jl | 6 +----- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/ext/SpeedyWeatherJLArraysExt.jl b/ext/SpeedyWeatherJLArraysExt.jl index c80745ec2..b3d3c099c 100644 --- a/ext/SpeedyWeatherJLArraysExt.jl +++ b/ext/SpeedyWeatherJLArraysExt.jl @@ -2,7 +2,9 @@ module SpeedyWeatherJLArraysExt using SpeedyWeather, JLArrays -# for RingGrids, every Array needs this method to strip away the parameters +# 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/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 From 7a96c92cdcbf796bd4f98277a1cda3452a3d1330 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Milan=20Kl=C3=B6wer?= Date: Thu, 25 Apr 2024 11:59:32 -0400 Subject: [PATCH 34/35] rm 2nd JLArrays in [extrax] in Project.toml --- Project.toml | 1 - 1 file changed, 1 deletion(-) diff --git a/Project.toml b/Project.toml index db1d73941..2ad23b26f 100644 --- a/Project.toml +++ b/Project.toml @@ -69,7 +69,6 @@ julia = "1.9" [extras] JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" -JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb" [targets] test = ["Test", "JLArrays"] From 1d5a1a74463cf31b8f5f6cf46a1ec9ac70155068 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Milan=20Kl=C3=B6wer?= Date: Thu, 25 Apr 2024 12:39:09 -0400 Subject: [PATCH 35/35] ::Type{<:JLArray} in JLArrays extension --- ext/SpeedyWeatherJLArraysExt.jl | 4 ++-- test/grids.jl | 2 -- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/ext/SpeedyWeatherJLArraysExt.jl b/ext/SpeedyWeatherJLArraysExt.jl index b3d3c099c..8f33c861e 100644 --- a/ext/SpeedyWeatherJLArraysExt.jl +++ b/ext/SpeedyWeatherJLArraysExt.jl @@ -4,7 +4,7 @@ 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 +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/test/grids.jl b/test/grids.jl index 8249b6529..69b7cad2a 100644 --- a/test/grids.jl +++ b/test/grids.jl @@ -370,8 +370,6 @@ end end end -RingGrids.nonparametric_type(::Type{<:JLArray}) = JLArray - @testset "AbstractGridArray: GPU (JLArrays)" begin NF = Float32 @testset for Grid in (