GeoRefDatasets

GeoRefDatasets.jl

"""
Example module for georeferenced datasets


"""
module GeoRefDatasets

using Test
import Base: getindex, size, view

"""
An array of size `sz`, for which all elements
            A[i₁,i₂,...,iₖ,...] = data[iₖ]
where `k` is the dimension to be kept.

It should be possible to extend it in the case where data has 2 or more
dimensions.
"""
mutable struct RepeatedArray{T,N,TD <: AbstractVector{T}} <: AbstractArray{T,N}
    sz::NTuple{N,Int}
    keepdim::Int
    data::TD
end

Base.size(ra::RepeatedArray) = ra.sz
Base.getindex(ra::RepeatedArray,I...) = ra.data[I[ra.keepdim]]


mutable struct GeoRefData
    # abstract array with data
    data
    # named tuple with coordinates
    # (lon = array of same size as data, lat, ...)
    grid
end

Base.getindex(d::GeoRefData,I...) = GeoRefData(
    d.data[I...],
    map(c -> view(c,I...),d.grid))

# return the underlying value
value(d::GeoRefData) = d.data
value(d::GeoRefData,i,J...) = d.data[i,J...]
# name can be exotic dimensions like :frequency or :ensemble_member
@inline coord(d::GeoRefData,name::Symbol,I...) = getproperty(d.grid,name)[I...]
#volumne(d::GeoRefData,name::Symbol,I...) = ...

# most of the time this would be enought: these names are quite standard
for cd in [:lon,:lat,:time,:elevation]
    @eval begin
        @inline ($cd)(d::GeoRefData,I...) = coord(d,Symbol($cd),I...)
    end
end

function subtest(data,gridlon,gridlat,gd)
    @test value(gd,2,3) == data[2,3]

    # slice a GeoRefDatasets
    sub_gd = gd[:,2]
    sub_sub_gd = sub_gd[3]
    @test value(sub_sub_gd) == data[3,2]

    # get all data in one go as an Array
    @test value(gd) == data

    # get the coordinates
    @test lon(gd,2,3) == gridlon[2,3]
    @test lat(gd,2,3) == gridlat[3,3]

    @test lon(gd,:,3) == gridlon[:,3]
    @test all(lon(gd,2,:) .== gridlon[2,:])

end

function test()
    dvec = [10,20,30]
    ra = GeoRefDatasets.RepeatedArray((10,2,3,4,5),3,dvec)
    @test ra[1,1,3,1,1] == dvec[3]


    sz = (100,101)
    glon = LinRange(-180,180,sz[1])
    glat = LinRange(-90,90,sz[2])

    # replicate lon/lat (which is not memory efficient)
    gridlon = repeat(glon, inner = (1,sz[2]));
    gridlat = repeat(glat',inner = (sz[1],1));

    data = gridlon .+ gridlat;

    gd = GeoRefData(data,(lon = gridlon, lat = gridlat))
    subtest(data,gridlon,gridlat,gd)

    # used RepeatedArray for lon/lat (which is memory efficient)
    gd = GeoRefData(data,(
        lon = RepeatedArray(sz,1,glon),
        lat = RepeatedArray(sz,2,glat)))

    subtest(data,gridlon,gridlat,gd)

end
export GeoRefData, lon, lat, time, elevation
end