implement read/write of .uns

Author: ilia-kats

src/anndata.jl

@@ -19,6 +19,8 @@ mutable struct AnnData <: AbstractAnnData
 
     layers::AbstractAlignedMapping{Tuple{1 => 1, 2 => 2}, String}
 
+    uns::Dict{<:AbstractString, <:Any}
+
     function AnnData(file::Union{HDF5.File, HDF5.Group}, backed=true, checkversion=true)
         if checkversion
             attrs = attributes(file)
@@ -71,6 +73,11 @@ mutable struct AnnData <: AbstractAnnData
             adata.layers = BackedAlignedMapping{Tuple{1 => 1, 2 => 2}}(adata, adata.file, "layers")
         end
 
+        # unstructured
+        adata.uns =
+            haskey(file, "uns") ? read_dict_of_mixed(file["uns"], separate_index=false) :
+            Dict{String, Any}()
+
         return adata
     end
 
@@ -91,6 +98,7 @@ mutable struct AnnData <: AbstractAnnData
         obsp::Union{AbstractDict{<:AbstractString, <:AbstractMatrix{<:Number}}, Nothing}=nothing,
         varp::Union{AbstractDict{<:AbstractString, <:AbstractMatrix{<:Number}}, Nothing}=nothing,
         layers::Union{AbstractDict{<:AbstractString, <:AbstractMatrix{<:Number}}, Nothing}=nothing,
+        uns::Union{AbstractDict{<:AbstractString, <:Any}, Nothing}=nothing,
     )
         m, n = size(X)
         if isnothing(obs)
@@ -122,6 +130,7 @@ mutable struct AnnData <: AbstractAnnData
         adata.varm = StrAlignedMapping{Tuple{1 => 2}}(adata, varm)
         adata.varp = StrAlignedMapping{Tuple{1 => 2, 2 => 2}}(adata, varp)
         adata.layers = StrAlignedMapping{Tuple{1 => 1, 2 => 2}}(adata, layers)
+        adata.uns = isnothing(uns) ? Dict{String, Any}() : uns
 
         return adata
     end
@@ -190,24 +199,25 @@ function Base.write(parent::Union{HDF5.File, HDF5.Group}, adata::AbstractAnnData
     else
         write_attr(parent, "X", adata.X)
         write_attr(parent, "layers", adata.layers)
-        write_unbacked(parent, adata)
+        write_metadata(parent, adata)
     end
 end
 
 function Base.write(adata)
     if file(adata) === nothing
-        throw("adata is not backed, need somewhere to write to")
+        error("adata is not backed, need somewhere to write to")
     end
-    write_unbacked(file(adata), adata)
+    write_metadata(file(adata), adata)
 end
 
-function write_unbacked(parent::Union{HDF5.File, HDF5.Group}, adata::AbstractAnnData)
+function write_metadata(parent::Union{HDF5.File, HDF5.Group}, adata::AbstractAnnData)
     write_attr(parent, "obs", adata.obs, index=adata.obs_names)
     write_attr(parent, "obsm", adata.obsm, index=adata.obs_names)
     write_attr(parent, "obsp", adata.obsp)
     write_attr(parent, "var", adata.var, index=adata.var_names)
     write_attr(parent, "varm", adata.varm, index=adata.var_names)
     write_attr(parent, "varp", adata.varp)
+    write_attr(parent, "uns", adata.uns)
 end
 # FileIO support
 load(f::File{format"h5ad"}) = readh5ad(filename(f), backed=false) # I suppose this is more consistent with the rest of FileIO?

src/hdf5_io.jl

@@ -1,13 +1,7 @@
-function read_dataframe(tablegroup::HDF5.Group)
+function read_dataframe(tablegroup::HDF5.Group; separate_index=true, kwargs...)
     columns = read_attribute(tablegroup, "column-order")
 
-    havecat = false
-    if haskey(tablegroup, "__categories")
-        havecat = true
-        catcols = tablegroup["__categories"]
-    end
-
-    if haskey(attributes(tablegroup), "_index")
+    if separate_index && haskey(attributes(tablegroup), "_index")
         indexdsetname = read_attribute(tablegroup, "_index")
         rownames = read(tablegroup[indexdsetname])
     else
@@ -17,26 +11,42 @@ function read_dataframe(tablegroup::HDF5.Group)
     df = DataFrame()
 
     for col in columns
-        column = read(tablegroup, col)
-        if havecat && haskey(catcols, col)
-            cats = read(catcols, col)
-            column = compress(CategoricalArray(map(x -> cats[x + 1], column)))
+        column = read_matrix(tablegroup[col])
+        if sum(size(column) .> 1) > 1
+            @warn "column $col has more than 1 dimension for data frame $(HDF5.name(tablegroup)), skipping"
         end
         df[!, col] = column
     end
 
     return df, rownames
 end
 
-function read_matrix(f::HDF5.Dataset)
+function read_matrix(f::HDF5.Dataset; kwargs...)
     mat = read(f)
-    if ndims(mat) > 1
+    if ndims(f) == 0
+        return mat
+    end
+    if ndims(f) > 1
         mat = PermutedDimsArray(mat, ndims(mat):-1:1) # transpose for h5py compatibility
     end
+    if haskey(attributes(f), "categories")
+        categories = f[read_attribute(f, "categories")]
+        ordered =
+            haskey(attributes(categories), "ordered") &&
+            read_attribute(categories, "ordered") == true
+        cats = read(categories)
+        mat = mat .+ 0x1
+        mat = compress(
+            CategoricalArray{eltype(cats), ndims(mat)}(
+                mat,
+                CategoricalPool{eltype(cats), eltype(mat)}(cats, ordered),
+            ),
+        )
+    end
     return mat
 end
 
-function read_matrix(f::HDF5.Group)
+function read_matrix(f::HDF5.Group; kwargs...)
     enctype = read_attribute(f, "encoding-type")
 
     if enctype == "csc_matrix" || enctype == "csr_matrix"
@@ -61,30 +71,46 @@ function read_matrix(f::HDF5.Group)
         mat = SparseMatrixCSC(shape..., indptr, indices, data)
         return iscsr ? mat' : mat
     else
-        throw("unknown encoding $enctype")
+        error("unknown encoding $enctype")
     end
 end
 
-function read_dict_of_matrices(f::HDF5.Group)
-    return Dict{String, AbstractArray{<:Number}}(key => read_matrix(f[key]) for key in keys(f))
+function read_dict_of_matrices(f::HDF5.Group; kwargs...)
+    return Dict{String, AbstractArray{<:Number}}(
+        key => read_matrix(f[key]; kwargs...) for key in keys(f)
+    )
 end
 
-read_auto(f::HDF5.Dataset) = (read_matrix(f), nothing)
-function read_auto(f::HDF5.Group)
-    enctype = read_attribute(f, "encoding-type")
-    if enctype == "dataframe"
-        return read_dataframe(f)
-    elseif endswith(enctype, "matrix")
-        return read_matrix(f), nothing
+read_auto(f::HDF5.Dataset; kwargs...) = (read_matrix(f; kwargs...), nothing)
+function read_auto(f::HDF5.Group; kwargs...)
+    if haskey(attributes(f), "encoding-type")
+        enctype = read_attribute(f, "encoding-type")
+        if enctype == "dataframe"
+            return read_dataframe(f; kwargs...)
+        elseif endswith(enctype, "matrix")
+            return read_matrix(f; kwargs), nothing
+        else
+            error("unknown encoding $enctype")
+        end
     else
-        throw("unknown encoding $enctype")
+        return read_dict_of_mixed(f; kwargs...), nothing
     end
 end
 
-function read_dict_of_mixed(f::HDF5.Group)
-    ret = Dict{String, Union{DataFrame, AbstractArray{<:Number}}}()
+function read_dict_of_mixed(f::HDF5.Group; kwargs...)
+    ret = Dict{
+        String,
+        Union{
+            DataFrame,
+            <:AbstractArray{<:Number},
+            <:AbstractArray{<:AbstractString},
+            <:AbstractString,
+            <:Number,
+            Dict,
+        },
+    }()
     for k in keys(f)
-        ret[k] = read_auto(f[k])[1] # assume data frames are properly aligned, so we can discard rownames
+        ret[k] = read_auto(f[k]; kwargs...)[1] # assume data frames are properly aligned, so we can discard rownames
     end
     return ret
 end
@@ -96,6 +122,15 @@ function write_attr(parent::Union{HDF5.File, HDF5.Group}, name::AbstractString,
     write_impl(parent, name, data; kwargs...)
 end
 
+function write_impl(
+    parent::Union{HDF5.File, HDF5.Group},
+    name::AbstractString,
+    data::Union{<:Number, <:AbstractString};
+    kwargs...,
+)
+    parent[name] = data
+end
+
 function write_impl(
     parent::Union{HDF5.File, HDF5.Group},
     name::AbstractString,
@@ -113,19 +148,20 @@ end
 function write_impl(
     parent::Union{HDF5.File, HDF5.Group},
     name::AbstractString,
-    data::CategoricalVector;
+    data::CategoricalArray;
     kwargs...,
 )
     write_impl(parent, name, data.refs .- 0x1; kwargs...)
     write_impl(parent, "__categories/$name", levels(data); kwargs...)
+    attributes(parent["__categories/$name"])["ordered"] = UInt8(isordered(data))
     attributes(parent[name])["categories"] = HDF5.Reference(parent["__categories"], name)
 end
 
 function write_impl(
     parent::Union{HDF5.File, HDF5.Group},
     name::AbstractString,
     data::AbstractDataFrame;
-    index::AbstractVector{<:AbstractString},
+    index::AbstractVector{<:AbstractString}=nothing,
     kwargs...,
 )
     g = create_group(parent, name)
@@ -140,8 +176,18 @@ function write_impl(
 
     idxname = "_index"
     columns = names(data)
-    while idxname ∈ columns
-        idxname = "_" * idxname
+    if !isnothing(index)
+        while idxname ∈ columns
+            idxname = "_" * idxname
+        end
+    else
+        if idxname ∈ columns
+            index = data[!, idxname]
+            select!(data, Not(idxname))
+        else
+            @warn "Data frame $(HDF5.name(parent))/$name does not have an _index column, a row number index will be written"
+            index = 1:nrow(data)
+        end
     end
     g = parent[name]
     write_impl(g, idxname, values(index); kwargs...)
@@ -157,22 +203,16 @@ write_impl(
     extensible::Bool=false,
     compress::UInt8=UInt8(9),
     kwargs...,
-) = write_impl(
-    parent,
-    name,
-    Int8.(data);
-    extensible=extensible,
-    compress=compress,
-    kwargs...,
-)
+) = write_impl(parent, name, Int8.(data); extensible=extensible, compress=compress, kwargs...)
 
 write_impl(parent::Union{HDF5.File, HDF5.Group}, name::AbstractString, data::SubArray; kwargs...) =
     write_impl(parent, name, copy(data); kwargs...)
 
 function write_impl(
     parentgrp::Union{HDF5.File, HDF5.Group},
     name::AbstractString,
-    data::AbstractArray,;
+    data::AbstractArray,
+    ;
     extensible::Bool=false,
     compress::UInt8=0x9,
     kwargs...,

src/mudata.jl

@@ -16,6 +16,8 @@ mutable struct MuData <: AbstractMuData
     varp::StrAlignedMapping{Tuple{1 => 2, 2 => 2}, MuData}
     varmap::StrAlignedMapping{Tuple{1 => 2}, MuData}
 
+    uns::Dict{<:AbstractString, <:Any}
+
     function MuData(file::Union{HDF5.File, HDF5.Group}, backed=true, checkversion=true)
         if checkversion
             attrs = attributes(file)
@@ -62,6 +64,11 @@ mutable struct MuData <: AbstractMuData
             haskey(file, "varmap") ? read_dict_of_matrices(file["varmap"]) : nothing,
         )
 
+        # unstructured
+        mdata.uns =
+            haskey(file, "uns") ? read_dict_of_mixed(file["uns"], separate_index=false) :
+            Dict{String, Any}()
+
         # Modalities
         mdata.mod = Dict{String, AnnData}()
         mods = HDF5.keys(file["mod"])
@@ -90,6 +97,7 @@ mutable struct MuData <: AbstractMuData
         varp::Union{AbstractDict{<:AbstractString, <:AbstractMatrix{<:Number}}, Nothing}=nothing,
         obsmap::Union{AbstractDict{<:AbstractString, <:AbstractVector{<:Integer}}, Nothing}=nothing,
         varmap::Union{AbstractDict{<:AbstractString, <:AbstractVector{<:Integer}}, Nothing}=nothing,
+        uns::Union{AbstractDict{<:AbstractString, <:Any}, Nothing}=nothing,
         do_update=true,
     )
         mdata = new(nothing, Dict{String, AnnData}())
@@ -108,6 +116,7 @@ mutable struct MuData <: AbstractMuData
         mdata.varp = StrAlignedMapping{Tuple{1 => 2, 2 => 2}}(mdata, varp)
         mdata.obsmap = StrAlignedMapping{Tuple{1 => 1}}(mdata, obsmap)
         mdata.varmap = StrAlignedMapping{Tuple{1 => 2}}(mdata, varmap)
+        mdata.uns = isnothing(uns) ? Dict{String, Any}() : uns
 
         if do_update
             update!(mdata)
@@ -197,11 +206,12 @@ function write_metadata(parent::Union{HDF5.File, HDF5.Group}, mudata::AbstractMu
     write_attr(parent, "obs", shrink_attr(mudata, :obs), index=mudata.obs_names)
     write_attr(parent, "obsm", mudata.obsm, index=mudata.obs_names)
     write_attr(parent, "obsp", mudata.obsp)
-    write_attr(parent, "obsmap", mudata.obsmap,)
+    write_attr(parent, "obsmap", mudata.obsmap)
     write_attr(parent, "var", shrink_attr(mudata, :var), index=mudata.var_names)
     write_attr(parent, "varm", mudata.varm, index=mudata.var_names)
     write_attr(parent, "varp", mudata.varp)
     write_attr(parent, "varmap", mudata.varmap)
+    write_attr(parent, "uns", mudata.uns)
 end
 
 # FileIO support