Merge pull request #7 from meggart/gdalmissings

Fix attributes for bands

Author: meggart

Project.toml

@@ -1,7 +1,7 @@
 name = "YAXArrayBase"
 uuid = "90b8fcef-0c2d-428d-9c56-5f86629e9d14"
 authors = ["Fabian Gans <[email protected]>"]
-version = "0.2.2"
+version = "0.3.0"
 
 [deps]
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"

src/YAXArrayBase.jl

@@ -31,7 +31,11 @@ function __init__()
 
   @require AxisIndices="f52c9ee2-1b1c-4fd8-8546-6350938c7f11" include("axisarrays/axisindices.jl")
 
-  @require ArchGDAL="c9ce4bd3-c3d5-55b8-8973-c0e20141b8c3" include("axisarrays/archgdal.jl")
+  @require ArchGDAL="c9ce4bd3-c3d5-55b8-8973-c0e20141b8c3" begin 
+    include("axisarrays/archgdal.jl")
+    include("datasets/archgdal.jl")
+  end
+
 
   @require Zarr="0a941bbe-ad1d-11e8-39d9-ab76183a1d99" include("datasets/zarr.jl")
 

src/axisarrays/archgdal.jl

@@ -1,6 +1,7 @@
 import .ArchGDAL: RasterDataset, AbstractRasterBand,
   getgeotransform, width, height, getname, getcolorinterp,
   getband, nraster, getdataset
+
 function dimname(a::RasterDataset, i)
     if i == 1
         return :Y
@@ -30,8 +31,17 @@ function dimvals(a::RasterDataset, i)
     end
 end
 iscontdim(a::RasterDataset, i) = i < 3 ? true : nraster(a)<8
-getattributes(a::RasterDataset) =
-  Dict{String,Any}("projection"=>ArchGDAL.toPROJ4(ArchGDAL.newspatialref(ArchGDAL.getproj(a))))
+function getattributes(a::RasterDataset)
+    globatts = Dict{String,Any}(
+        "projection_PROJ4"=>ArchGDAL.toPROJ4(ArchGDAL.newspatialref(ArchGDAL.getproj(a))),
+        "projection_WKT"=>ArchGDAL.toWKT(ArchGDAL.newspatialref(ArchGDAL.getproj(a))),
+    )
+    bands = (getbandattributes(ArchGDAL.getband(a, i)) for i in 1:size(a, 3))
+    allbands = mergewith(bands...) do a1,a2
+        isequal(a1,a2) ? a1 : missing
+    end
+    merge(globatts, allbands)
+end
 
 
 function dimname(::AbstractRasterBand, i)
@@ -54,5 +64,24 @@ function dimvals(b::AbstractRasterBand, i)
     end
 end
 iscontdim(a::AbstractRasterBand, i) = true
-getattributes(a::AbstractRasterBand) =
-  getattributes(ArchGDAL.RasterDataset(ArchGDAL.getdataset(a)))
+function getattributes(a::AbstractRasterBand)
+  atts = getattributes(ArchGDAL.RasterDataset(ArchGDAL.getdataset(a)))
+  bandatts = getbandattributes(a)
+  merge(atts, bandatts)
+end
+
+function insertattifnot!(attrs, val, name, condition)
+    if !condition(val)
+        attrs[name] = val
+    end
+end
+function getbandattributes(a::AbstractRasterBand)
+  atts = Dict{String,Any}()
+  catdict = Dict((i-1)=>v for (i,v) in enumerate(AG.getcategorynames(a)))
+  insertattifnot!(atts, AG.getnodatavalue(a), "missing_value", isnothing)
+  insertattifnot!(atts, catdict, "labels", isempty)
+  insertattifnot!(atts, AG.getunittype(a), "units", isempty)
+  insertattifnot!(atts, AG.getoffset(a), "add_offset", iszero)
+  insertattifnot!(atts, AG.getscale(a), "scale_factor", x->isequal(x, one(x)))
+  atts
+end

src/datasets/archgdal.jl

@@ -0,0 +1,218 @@
+import .ArchGDAL: RasterDataset, AbstractRasterBand,
+  getgeotransform, width, height, getname, getcolorinterp,
+  getband, nraster, getdataset, ArchGDAL
+using .ArchGDAL.DiskArrays: GridChunks, DiskArrays, eachchunk
+const AG = ArchGDAL
+
+struct GDALBand{T} <: AG.DiskArrays.AbstractDiskArray{T,2}
+    filename::String
+    band::Int
+    size::Tuple{Int,Int}
+    attrs::Dict{String,Any}
+    cs::GridChunks{2}
+end
+function GDALBand(b, filename, i)
+    s = size(b)
+    atts = getbandattributes(b)
+    GDALBand{AG.pixeltype(b)}(filename, i, s, atts, eachchunk(b))
+end
+Base.size(b::GDALBand) = b.size
+DiskArrays.eachchunk(b::GDALBand) = b.cs
+DiskArrays.haschunks(::GDALBand) = DiskArrays.Chunked()
+function DiskArrays.readblock!(b::GDALBand, aout, r::AbstractUnitRange...)
+    AG.read(b.filename) do ds
+        AG.getband(ds, b.band) do bh
+            DiskArrays.readblock!(bh, aout, r...)
+        end
+    end
+end
+function DiskArrays.writeblock!(b::GDALBand, ain, r::AbstractUnitRange...)
+    AG.read(b.filename, flags=AG.OF_Update) do ds
+        AG.getband(ds, b.band) do bh
+            DiskArrays.writeblock!(bh, ain, r...)
+        end
+    end
+end
+
+struct GDALDataset
+    filename::String
+    bandsize::Tuple{Int,Int}
+    projection::String
+    trans::Vector{Float64}
+    bands::OrderedDict{String}
+end
+
+function GDALDataset(filename)
+    AG.read(filename) do r
+        nb = AG.nraster(r)
+        allbands = map(1:nb) do iband
+            b = AG.getband(r, iband)
+            gb = GDALBand(b, filename, iband)
+            name = AG.GDAL.gdalgetdescription(b.ptr)
+            if isempty(name)
+                name = AG.getname(AG.getcolorinterp(b))
+            end
+            name => gb
+        end
+        proj = AG.getproj(r)
+        trans = AG.getgeotransform(r)
+        s = AG._common_size(r)
+        allnames = first.(allbands)
+        if !allunique(allnames)
+            allbands = ["Band$i"=>last(v) for (i,v) in enumerate(allbands)]
+        end
+        GDALDataset(filename, s[1:end-1], proj, trans, OrderedDict(allbands))
+    end
+end
+Base.haskey(ds::GDALDataset, k) = in(k, ("X", "Y")) || haskey(ds.bands, k)
+#Implement Dataset interface
+function get_var_handle(ds::GDALDataset, name)
+    if name == "X"
+        range(ds.trans[1], length = ds.bandsize[1], step = ds.trans[2])
+    elseif name == "Y"
+        range(ds.trans[4], length = ds.bandsize[2], step = ds.trans[6])
+    else
+        ds.bands[name]
+    end
+end
+
+
+get_varnames(ds::GDALDataset) = collect(keys(ds.bands))
+
+function get_var_dims(ds::GDALDataset, d) 
+    if d === "X"
+        return ("X",)
+    elseif d==="Y"
+        return ("Y",)
+    else
+        return ("X", "Y")
+    end
+end
+
+get_global_attrs(ds::GDALDataset) = Dict("projection"=>ds.projection)
+
+function get_var_attrs(ds::GDALDataset, name)
+    if name in ("Y", "X")
+        Dict{String,Any}()
+    else
+        ds.bands[name].attrs
+    end
+end
+
+const colornames = AG.getname.(AG.GDALColorInterp.(0:16))
+
+islat(s) = startswith(uppercase(s), "LAT")
+islon(s) = startswith(uppercase(s), "LON")
+isx(s) = uppercase(s) == "X"
+isy(s) = uppercase(s) == "Y"
+
+function totransform(x, y)
+    xstep = diff(x)
+    ystep = diff(y)
+    if !all(isapprox(first(xstep)), xstep) || !all(isapprox(first(ystep)), ystep)
+        throw(ArgumentError("Grid must have regular spacing"))
+    end
+    Float64[first(x), first(xstep), 0.0, first(y), 0.0, first(ystep)]
+end
+totransform(x::AbstractRange, y::AbstractRange) =
+    Float64[first(x), step(x), 0.0, first(y), 0.0, step(y)]
+getproj(userproj::String, attrs) = AG.importPROJ4(userproj)
+getproj(userproj::AG.AbstractSpatialRef, attrs) = userproj
+function getproj(userproj::Nothing, attrs)
+    if haskey(attr, "projection_PROJ4")
+        return AG.importPROJ4(attr["projection_PROJ4"])
+    elseif haskey(attr, "projection_WKT")
+        return AG.importWKT(attr["projection_WKT"])
+    else
+        error(
+            "Could not determine output projection from attributes, please specify userproj",
+        )
+    end
+end
+
+
+function create_dataset(
+    ::Type{<:GDALDataset},
+    outpath,
+    gatts,
+    dimnames,
+    dimvals,
+    dimattrs,
+    vartypes,
+    varnames,
+    vardims,
+    varattrs,
+    varchunks;
+    userproj = nothing,
+    kwargs...,
+)
+    @assert length(dimnames) == 2
+    proj, trans = if islon(dimnames[1]) && islat(dimnames[2])
+        #Lets set srs to EPSG:4326
+        proj = AG.importEPSG(4326)
+        trans = totransform(dimvals[1], dimvals[2])
+        proj, trans
+    elseif isx(dimnames[1]) && isy(dimnames[2])
+        #Try to find out srs
+        proj = getproj(userproj, gatts)
+        trans = totransform(dimvals[1], dimvals[2])
+        proj, trans
+    else
+        error("Did not find x, y or lon, lat dimensions in dataset")
+    end
+    cs = first(varchunks)
+    @assert all(isequal(varchunks[1]), varchunks)
+
+    driver = AG.getdriver(AG.extensiondriver(outpath))
+
+    nbands = length(varnames)
+    dtype = promote_type(vartypes...)
+    s = (length.(dimvals)...,)
+    bands = AG.create(
+        outpath;
+        driver = driver,
+        width = length(dimvals[1]),
+        height = length(dimvals[2]),
+        nbands = nbands,
+        dtype = dtype,
+        options = ["BLOCKXSIZE=$(cs[1])", "BLOCKYSIZE=$(cs[2])"]
+    ) do ds
+        AG.setgeotransform!(ds, trans)
+        bands = map(1:length(varnames)) do i
+            b = AG.getband(ds, i)
+            icol = findfirst(isequal(varnames[i]), colornames)
+            if isnothing(icol)
+                AG.setcolorinterp!(b, AG.GDAL.GDALColorInterp(0))
+            else
+                AG.setcolorinterp!(b, AG.GDAL.GDALColorInterp(icol - 1))
+            end
+            AG.GDAL.gdalsetdescription(b.ptr, varnames[i])
+            atts = varattrs[i]
+            haskey(atts, "missing_value") && AG.setnodatavalue!(b, atts["missing_value"])
+            if haskey(atts, "labels")
+                labeldict = atts[labels]
+                maxlabel = maximum(keys(labeldict))
+                kt = keytype(labeldict)
+                labelvec = [haskey(labeldict, et(i)) ? labeldict[et(i)] : "" for i = 0:maxlabel]
+                AG.setcategorynames!(b, labelvec)
+            end
+            haskey(atts, "units") && AG.setunittype!(b, atts["units"])
+            haskey(atts, "scale_factor") && AG.setscale!(b, atts["scale_factor"])
+            haskey(atts, "add_offset") && AG.setoffset!(b, atts["add_offset"])
+            GDALBand{dtype}(outpath, i, s, atts, AG.DiskArrays.GridChunks(s,cs))
+        end
+    end
+    return GDALDataset(outpath, s, AG.toPROJ4(proj), trans, OrderedDict(vn=>b for (vn,b) in zip(varnames, bands)))
+end
+
+allow_parallel_write(::Type{<:GDALDataset}) = false
+allow_parallel_write(::GDALDataset) = false
+
+allow_missings(::Type{<:GDALDataset}) = false
+allow_missings(::GDALDataset) = false
+
+backendlist[:gdal] = GDALDataset
+push!(backendregex,r".tif$"=>GDALDataset)
+push!(backendregex,r".gtif$"=>GDALDataset)
+push!(backendregex,r".tiff$"=>GDALDataset)
+push!(backendregex,r".gtiff$"=>GDALDataset)

src/datasets/datasetinterface.jl

@@ -11,11 +11,12 @@ function get_var_dims(ds, name) end
 "Return a dict with the attributes for a given variable"
 function get_var_attrs(ds,name) end
 
+"Return a dict with global attributes for the dataset"
+function get_global_attrs(ds) end
 
 #Functions to be implemented for Dataset sinks
 "Initialize and return a handle to a new empty dataset"
-create_empty(T::Type,path) =
-  error("create_empty not implemented for $T")
+function create_empty(T::Type,path,gatts) end
 
 """
     add_var(ds, T, name, s, dimlist, atts)
@@ -41,13 +42,28 @@ Returns true if Union{T,Missing} is an allowed data type for the backend
 and if an array containing missings can be written to the array.
 """
 allow_missings(ds) = true
+
 #Fallback for writing array
 function add_var(ds,x::AbstractArray,name,dimlist,atts;kwargs...)
   a = add_var(ds,eltype(x),name,size(x),dimlist,atts;kwargs...)
   a .= x
   a
 end
 
+function create_dataset(T::Type, path, gatts, dimnames, dimvals, dimattrs, vartypes, varnames, vardims, varattrs, varchunks; kwargs...)
+  ds = create_empty(T, path, gatts)
+  axlengths = Dict{String, Int}()
+  for (dname, dval, dattr) in zip(dimnames, dimvals, dimattrs)
+    add_var(ds, dval, dname, (dname,), dattr)
+    axlengths[dname] = length(dval)
+  end
+  for (T, vn, vd, va, vc) in zip(vartypes, varnames, vardims, varattrs, varchunks)
+    s = getindex.(Ref(axlengths),vd) 
+    add_var(ds, T, vn, (s...,), vd, va; chunksize = vc, kwargs...)
+  end
+  ds
+end
+
 using DataStructures: OrderedDict
 """
     YAXArrayBase.backendlist

src/datasets/netcdf.jl

@@ -38,6 +38,7 @@ Base.size(v::NetCDFVariable) = v.size
 get_var_dims(ds::NetCDFDataset,name) = NetCDF.open(v->map(i->i.name,v[name].dim),ds.filename)
 get_varnames(ds::NetCDFDataset) = NetCDF.open(v->collect(keys(v.vars)),ds.filename)
 get_var_attrs(ds::NetCDFDataset, name) = NetCDF.open(v->v[name].atts,ds.filename)
+get_global_attrs(ds::NetCDFDataset) = NetCDF.open(nc->nc.gatts, ds.filename)
 function Base.getindex(ds::NetCDFDataset, i)
   s,et = NetCDF.open(j->(size(j),eltype(j)),ds.filename,i)
   NetCDFVariable{et,length(s)}(ds.filename, i, s)
@@ -51,13 +52,15 @@ function add_var(p::NetCDFDataset, T::Type, varname, s, dimnames, attr;
   NetCDFVariable{T,length(s)}(p.filename,varname,(s...,))
 end
 
-function create_empty(::Type{NetCDFDataset}, path)
-  NetCDF.create(_->nothing, path, NcVar[])
+function create_empty(::Type{NetCDFDataset}, path, gatts=Dict())
+  NetCDF.create(_->nothing, path, NcVar[], gatts = gatts)
   NetCDFDataset(path)
 end
 
+allow_parallel_write(::Type{<:NetCDFDataset}) = false
 allow_parallel_write(::NetCDFDataset) = false
 
+allow_missings(::Type{<:NetCDFDataset}) = false
 allow_missings(::NetCDFDataset) = false
 
 backendlist[:netcdf] = NetCDFDataset

src/datasets/zarr.jl

@@ -9,6 +9,7 @@ ZarrDataset(g::String) = ZarrDataset(zopen(g))
 get_var_dims(ds::ZarrDataset,name) = reverse(ds[name].attrs["_ARRAY_DIMENSIONS"])
 get_varnames(ds::ZarrDataset) = collect(keys(ds.g.arrays))
 get_var_attrs(ds::ZarrDataset, name) = ds[name].attrs
+get_global_attrs(ds::ZarrDataset) = ds.g.attrs
 Base.getindex(ds::ZarrDataset, i) = ds.g[i]
 Base.haskey(ds::ZarrDataset,k) = haskey(ds.g,k)
 
@@ -33,7 +34,7 @@ function add_var(p::ZarrDataset, a::AbstractArray, varname, dimnames, attr;
   a
 end
 
-create_empty(::Type{ZarrDataset}, path) = ZarrDataset(zgroup(path))
+create_empty(::Type{ZarrDataset}, path, gatts=Dict()) = ZarrDataset(zgroup(path, attrs=gatts))
 
 backendlist[:zarr] = ZarrDataset
 push!(backendregex, r"(.zarr$)|(.zarr/$)"=>ZarrDataset)