Add archgdal dataset driver

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

@@ -32,6 +32,9 @@ 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" include("datasets/archgdal.jl")
+
 
   @require Zarr="0a941bbe-ad1d-11e8-39d9-ab76183a1d99" include("datasets/zarr.jl")
 

src/axisarrays/archgdal.jl

@@ -1,7 +1,6 @@
 import .ArchGDAL: RasterDataset, AbstractRasterBand,
   getgeotransform, width, height, getname, getcolorinterp,
   getband, nraster, getdataset
-const AG = ArchGDAL
 
 function dimname(a::RasterDataset, i)
     if i == 1
@@ -33,7 +32,10 @@ function dimvals(a::RasterDataset, i)
 end
 iscontdim(a::RasterDataset, i) = i < 3 ? true : nraster(a)<8
 function getattributes(a::RasterDataset)
-    globatts = Dict{String,Any}("projection"=>ArchGDAL.toPROJ4(ArchGDAL.newspatialref(ArchGDAL.getproj(a))))
+    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
@@ -75,9 +77,10 @@ function insertattifnot!(attrs, val, name, condition)
 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, AG.getcategorynames(a), "Categories", isempty)
-  insertattifnot!(atts, AG.GDAL.gdalgetrasterunittype(a.ptr), "Units", isempty)
+  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

src/datasets/archgdal.jl

@@ -1,24 +1,35 @@
+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} <: AbstractDiskArray{T,2}
+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)
+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...)
+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...)
+        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
@@ -27,47 +38,167 @@ struct GDALDataset
     filename::String
     bandsize::Tuple{Int,Int}
     projection::String
-    trans::NTuple{6,Float64}
+    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)
+            b = AG.getband(r, iband)
+            gb = GDALBand(b, filename, iband)
             name = AG.getname(AG.getcolorinterp(b))
-            isempty(name) && (name = string("Band",iband))
             name => gb
         end
         proj = AG.getproj(r)
-        trans = (AG.getgeotransform(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)
+Base.haskey(ds::GDALDataset, k) = in(k, ("X", "Y")) || haskey(ds.bands, k)
 #Implement Dataset interface
-function get_var_handle(ds::GDALDataset, name) 
+function get_var_handle(ds::GDALDataset, name)
     if name == "X"
-        range(ds.trans[1],length=ds.bandsize[1], step=ds.trans[2])
+        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])
+        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))
 
 get_var_dims(ds::GDALDataset, _) = ("X", "Y")
 
-function get_var_attrs(ds::GDALDataset,name) 
-    if name in ("Y","X")
+function get_var_attrs(ds::GDALDataset, name)
+    if name in ("Y", "X")
         Dict{String,Any}()
     else
         ds.bands[name].attrs
     end
-end
+end
+
+const colornames = AG.getname.(AG.GDAL.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
+            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,6 +11,8 @@ 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"
@@ -49,16 +51,16 @@ function add_var(ds,x::AbstractArray,name,dimlist,atts;kwargs...)
   a
 end
 
-function create_dataset(T::Type, path, dimnames, dimvals, dimattrs, vartypes, varnames, vardims, varattrs, varchunks, kwargs)
-  ds = create_empty(T, path)
+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...)
+    add_var(ds, T, vn, (s...,), vd, va; chunksize = vc, kwargs...)
   end
   ds
 end

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,8 +52,8 @@ 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)
+  NetCDF.create(_->nothing, path, NcVar[], gatts = gatts)
   NetCDFDataset(path)
 end
 

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) = ZarrDataset(zgroup(path, attrs=gatts))
 
 backendlist[:zarr] = ZarrDataset
 push!(backendregex, r"(.zarr$)|(.zarr/$)"=>ZarrDataset)