Skip to content

Instantly share code, notes, and snippets.

@Alexander-Barth
Last active October 10, 2019 14:18
Show Gist options
  • Save Alexander-Barth/41f70c292ec9cdac2543e6c6a01ef612 to your computer and use it in GitHub Desktop.
Save Alexander-Barth/41f70c292ec9cdac2543e6c6a01ef612 to your computer and use it in GitHub Desktop.
GeoRefDatasets
"""
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment