views improvements

- obsmap/varmap now refer to the AnnDataViews in the view, not the
  parent object
- the modalities dict in a MuDataView is now ordered (by using freeze()
  from OrderedCollections.jl)
- various minor fixes

Author: ilia-kats

Project.toml

@@ -6,7 +6,6 @@ version = "0.2.1"
 [deps]
 CategoricalArrays = "324d7699-5711-5eae-9e2f-1d82baa6b597"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
-CompressHashDisplace = "33dc01b2-0ff1-4271-8670-ed883dd4b55b"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 EllipsisNotation = "da5c29d0-fa7d-589e-88eb-ea29b0a81949"
@@ -15,6 +14,7 @@ FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
+OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 PooledArrays = "2dfb63ee-cc39-5dd5-95bd-886bf059d720"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -25,14 +25,14 @@ Zarr = "0a941bbe-ad1d-11e8-39d9-ab76183a1d99"
 [compat]
 CategoricalArrays = "0.9, 0.10, 1"
 Compat = "4.10.0"
-CompressHashDisplace = "0.1.2"
 DataFrames = "1.5"
 DataStructures = "0.18, 0.19"
 EllipsisNotation = "1.8.0"
 FileIO = "1.6"
 FillArrays = "1.13.0"
 HDF5 = "0.16 - 0.99, 1"
 Logging = "1.6.7"
+OrderedCollections = "1.6"
 PooledArrays = "1"
 StructArrays = "0.6.4, 0.7"
 Zarr = "0.9"

src/Muon.jl

@@ -13,7 +13,7 @@ using StructArrays
 using PooledArrays
 using FillArrays
 using DataStructures
-import CompressHashDisplace: FrozenDict
+using OrderedCollections
 using FileIO
 
 using Compat

src/anndata.jl

@@ -355,7 +355,7 @@ function Base.getindex(
 )
     @boundscheck checkbounds(adata, I, J)
     i, j = convertidx(I, adata.obs_names), convertidx(J, adata.var_names)
-    newad = AnnData(
+    @inbounds newad = AnnData(
         X=adata.X[i, j],
         obs=isempty(adata.obs) ? nothing : adata.obs[i, :],
         obs_names=adata.obs_names[i],
@@ -394,10 +394,14 @@ struct AnnDataView{Ti, Tj} <: AbstractAnnData
     uns::Dict{<:AbstractString, <:Any}
 end
 
-function Base.view(ad::AnnData, I, J)
+@inline function Base.view(
+    ad::AnnData,
+    I::Union{OrdinalRange, Colon, AbstractVector{<:Integer}, AbstractVector{<:AbstractString}, Integer, AbstractString},
+    J::Union{OrdinalRange, Colon, AbstractVector{<:Integer}, AbstractVector{<:AbstractString}, Integer, AbstractString},
+)
     @boundscheck checkbounds(ad, I, J)
     i, j = convertidx(I, ad.obs_names), convertidx(J, ad.var_names)
-    X = isbacked(ad) ? nothing : @view ad.X[i, j]
+    X = isbacked(ad) ? nothing : @inbounds view(ad.X, i, j)
 
     return AnnDataView(
         ad,
@@ -416,10 +420,14 @@ function Base.view(ad::AnnData, I, J)
         ad.uns,
     )
 end
-function Base.view(ad::AnnDataView, I, J)
+function Base.view(
+    ad::AnnDataView,
+    I::Union{OrdinalRange, Colon, AbstractVector{<:Integer}, AbstractVector{<:AbstractString}, Integer, AbstractString},
+    J::Union{OrdinalRange, Colon, AbstractVector{<:Integer}, AbstractVector{<:AbstractString}, Integer, AbstractString},
+)
     @boundscheck checkbounds(ad, I, J)
     i, j = Base.reindex(parentindices(ad), (convertidx(I, ad.obs_names), convertidx(J, ad.var_names)))
-    return view(parent(ad), i, j)
+    return @inbounds view(parent(ad), i, j)
 end
 
 function Base.getproperty(ad::AnnDataView, s::Symbol)

src/index.jl

@@ -270,7 +270,7 @@ end
 end
 @inline function Base.view(idx::SubIndex, I)
     @boundscheck checkbounds(idx, I)
-    return @inbounds view(parent(idx), Base.reindex((parentindices(idx),), (I,))[1])
+    return @inbounds view(parent(idx), Base.reindex((parentindices(idx),), (I,))...)
 end
 
 Base.copy(si::SubIndex) = Index(si)
@@ -373,46 +373,39 @@ function Base.getindex(si::SubIndex{T}, elem::T, ::Val{false}) where {T}
     end
     return res
 end
-Base.@propagate_inbounds function Base.getindex(si::SubIndex, i::Union{Integer, AbstractVector{<:Integer}})
+@inline function Base.getindex(si::SubIndex, i::Union{Integer, AbstractVector{<:Integer}})
     @boundscheck checkbounds(si, i)
-    return parent(si)[Base.reindex((parentindices(si),), (i,))[1]]
+    return @inbounds parent(si)[Base.reindex((parentindices(si),), (i,))...]
 end
-Base.@propagate_inbounds function Base.getindex(
-    si::SubIndex{T, V, Colon},
-    i::Union{Integer, AbstractVector{<:Integer}},
-) where {T, V}
+@inline function Base.getindex(si::SubIndex{T, V, Colon}, i::Union{Integer, AbstractVector{<:Integer}}) where {T, V}
     @boundscheck checkbounds(si, i)
-    return parent(si)[i]
+    return @inbounds parent(si)[i]
 end
 
-Base.@propagate_inbounds function Base.setindex!(si::SubIndex{T}, newval::T, i::Integer) where {T}
+@inline function Base.setindex!(si::SubIndex{T}, newval::T, i::Integer) where {T}
     @boundscheck checkbounds(si, i)
-    setindex!(parent(si), newval, Base.reindex((parentindices(si),), (i,))[1])
+    @inbounds setindex!(parent(si), newval, Base.reindex((parentindices(si),), (i,))...)
     return si
 end
-Base.@propagate_inbounds function Base.setindex!(si::SubIndex{T, V, Colon}, newval::T, i::Integer) where {T, V}
+@inline function Base.setindex!(si::SubIndex{T, V, Colon}, newval::T, i::Integer) where {T, V}
     @boundscheck checkbounds(si, i)
-    return setindex!(parent(si), newval, i)
+    return @inbounds setindex!(parent(si), newval, i)
 end
-Base.@propagate_inbounds function Base.setindex!(si::SubIndex{T}, newval::T, oldval::T) where {T}
+function Base.setindex!(si::SubIndex{T}, newval::T, oldval::T) where {T}
     oldidx = parent(si)[oldval, true]
     foldidx = findfirst(in(parentindices(si)), oldidx)
     if isnothing(foldidx)
         throw(KeyError(oldval))
     end
-    parent(si)[oldidx[foldidx]] = newval
+    @inbounds parent(si)[oldidx[foldidx]] = newval
     return si
 end
-Base.@propagate_inbounds function Base.setindex!(
-    si::SubIndex{T, V, I},
-    newval::T,
-    oldval::T,
-) where {T, V, I <: AbstractArray{<:Integer}}
+function Base.setindex!(si::SubIndex{T, V, I}, newval::T, oldval::T) where {T, V, I <: AbstractArray{<:Integer}}
     oldidx = parent(si)[oldval, true]
     foldidx = findfirst(in(si.revmapping), oldidx)
     if isnothing(foldidx)
         throw(KeyError(oldval))
     end
-    parent(si)[oldidx[foldidx]] = newval
+    @inbounds parent(si)[oldidx[foldidx]] = newval
     return si
 end

src/mudata.jl

@@ -361,18 +361,15 @@ Base.setindex!(mdata::MuData, ad::AnnData, key::AbstractString) = setindex!(mdat
 Base.setindex!(mdata::MuData, ad::AnnData, key::Symbol) = setindex!(mdata.mod, ad, string(key))
 
 function Base.getindex(
-    mdata::MuData,
+    mdata::AbstractMuData,
     I::Union{OrdinalRange, Colon, AbstractVector{<:Integer}, AbstractVector{<:AbstractString}, Number, AbstractString},
     J::Union{OrdinalRange, Colon, AbstractVector{<:Integer}, AbstractVector{<:AbstractString}, Number, AbstractString},
 )
     @boundscheck checkbounds(mdata, I, J)
     i, j = convertidx(I, mdata.obs_names), convertidx(J, mdata.var_names)
-    newmu = MuData(
+    @inbounds newmu = MuData(
         mod=OrderedDict{String, AnnData}(
-            k => ad[
-                getadidx(i, vec(mdata.obsmap[k]), mdata.obs_names),
-                getadidx(j, vec(mdata.varmap[k]), mdata.var_names),
-            ] for (k, ad) ∈ mdata.mod
+            k => ad[getadidx(i, vec(mdata.obsmap[k])), getadidx(j, vec(mdata.varmap[k]))] for (k, ad) ∈ mdata.mod
         ),
         obs=isempty(mdata.obs) ? nothing : mdata.obs[i, :],
         obs_names=mdata.obs_names[i],
@@ -394,13 +391,9 @@ function Base.getindex(
     return newmu
 end
 
-getadidx(I::Colon, ref::AbstractVector{<:Unsigned}, idx::AbstractIndex{<:AbstractString}, reduce_memory=false) = I
-function getadidx(
-    I::Union{AbstractVector{<:Integer}, OrdinalRange},
-    ref::AbstractVector{<:Unsigned},
-    idx::AbstractIndex{<:AbstractString},
-    reduce_memory=false,
-)
+getadidx(I::Colon, ref::AbstractVector{<:Unsigned}, reduce_memory=false) = I
+getadidx(I::OrdinalRange, ref::AbstractVector{<:Unsigned}, reduce_memory=false) = filter(>(0x0), ref[I])
+function getadidx(I::AbstractVector{<:Integer}, ref::AbstractVector{<:Unsigned}, reduce_memory=false)
     J = filter(>(0x0), ref[I])
     if reduce_memory && length(J) > 0
         diff = J[end] - J[1]
@@ -414,8 +407,7 @@ function getadidx(
     end
     return J
 end
-getadidx(I::Number, ref::AbstractVector{<:Unsigned}, idx::Index{<:AbstractString}, reduce_memory=false) =
-    getadidx([I], ref, idx, reduce_memory)
+getadidx(I::Number, ref::AbstractVector{<:Unsigned}, reduce_memory=false) = getadidx([I], ref, reduce_memory)
 
 adjustmap!(map::AbstractVector{<:Unsigned}, I::Colon) = map
 function adjustmap!(map::AbstractVector{<:Unsigned}, I::Union{Integer, AbstractVector{<:Integer}, OrdinalRange})
@@ -999,77 +991,78 @@ push_var!(
 ) where {N, M} =
     _push_attr!(mdata, :var, 0x2, columns, mods, common=common, prefixed=prefixed, drop=drop, only_drop=only_drop)
 
-struct MuDataView{Ti, Tj} <: AbstractMuData
+mutable struct MuDataView{Ti, Tj} <: AbstractMuData
     parent::MuData
     I::Ti
     J::Tj
 
-    mod::FrozenDict{String, <:AnnDataView}
+    mod::LittleDict{String, <:AnnDataView}
     obs::SubDataFrame
     obs_names::SubIndex{<:AbstractString}
     obsm::StrAlignedMappingView{Tuple{1 => 1}}
     obsp::StrAlignedMappingView{Tuple{1 => 1, 2 => 1}}
-    obsmap::StrAlignedMappingView{Tuple{1 => 1}}
+    obsmap::StrAlignedMapping{Tuple{1 => 1}, MuDataView{Ti, Tj}, false}
 
     var::SubDataFrame
     var_names::SubIndex{<:AbstractString}
     varm::StrAlignedMappingView{Tuple{1 => 2}}
     varp::StrAlignedMappingView{Tuple{1 => 2, 2 => 2}}
-    varmap::StrAlignedMappingView{Tuple{1 => 2}}
+    varmap::StrAlignedMapping{Tuple{1 => 2}, MuDataView{Ti, Tj}, false}
 
     uns::Dict{<:AbstractString, <:Any}
+
+    function MuDataView{Ti, Tj}(parent::MuData, I::Ti, J::Tj) where {Ti, Tj}
+        mv = new(
+            parent,
+            I,
+            J,
+            freeze(
+                OrderedDict(
+                    m => view(ad, getadidx(I, vec(parent.obsmap[m]), true), getadidx(J, vec(parent.varmap[m]), true)) for (m, ad) ∈ parent.mod
+                ),
+            ),
+        )
+
+        mv.obs = view(parent.obs, nrow(parent.obs) > 0 ? I : (:), :)
+        mv.obs_names = view(parent.obs_names, I)
+        mv.obsm = view(parent.obsm, I)
+        mv.obsp = view(parent.obsp, I, I)
+
+        mv.var = view(parent.var, nrow(parent.var) > 0 ? J : (:), :)
+        mv.var_names = view(parent.var_names, J)
+        mv.varm = view(parent.varm, J)
+        mv.varp = view(parent.varp, J, J)
+
+        mv.obsmap = StrAlignedMapping{Tuple{1 => 1}, false}(mv)
+        mv.varmap = StrAlignedMapping{Tuple{1 => 2}, false}(mv)
+        copy_subset(parent.obsmap, mv.obsmap, I, J),
+        copy_subset(parent.varmap, mv.varmap, I, J),
+        for mod ∈ keys(mv.mod)
+            adjustmap!(vec(mv.obsmap[mod]), I)
+            adjustmap!(vec(mv.varmap[mod]), J)
+        end
+        return mv
+    end
 end
+MuDataView(parent::MuData, I::Ti, J::Tj) where {Ti, Tj} = MuDataView{Ti, Tj}(parent, I, J)
 
-function Base.view(
+@inline function Base.view(
     mu::MuData,
-    I::Union{OrdinalRange, Colon, AbstractVector{<:Integer}, AbstractVector{<:AbstractString}, Number, AbstractString},
-    J::Union{OrdinalRange, Colon, AbstractVector{<:Integer}, AbstractVector{<:AbstractString}, Number, AbstractString},
+    I::Union{OrdinalRange, Colon, AbstractVector{<:Integer}, AbstractVector{<:AbstractString}, Integer, AbstractString},
+    J::Union{OrdinalRange, Colon, AbstractVector{<:Integer}, AbstractVector{<:AbstractString}, Integer, AbstractString},
 )
     @boundscheck checkbounds(mu, I, J)
     i, j = convertidx(I, mu.obs_names), convertidx(J, mu.var_names)
-    mod = Dict(
-        m => view(
-            ad,
-            getadidx(i, vec(mu.obsmap[m]), mu.obs_names, true),
-            getadidx(j, vec(mu.varmap[m]), mu.var_names, true),
-        ) for (m, ad) ∈ mu.mod
-    )
-    return MuDataView(
-        mu,
-        i,
-        j,
-        FrozenDict(mod),
-        view(mu.obs, nrow(mu.obs) > 0 ? i : (:), :),
-        view(mu.obs_names, i),
-        view(mu.obsm, i),
-        view(mu.obsp, i, i),
-        view(mu.obsmap, i),
-        view(mu.var, nrow(mu.var) > 0 ? j : (:), :),
-        view(mu.var_names, j),
-        view(mu.varm, j),
-        view(mu.varp, j, j),
-        view(mu.varmap, j),
-        mu.uns,
-    )
+    return MuDataView(mu, i, j)
 end
 function Base.view(
     mu::MuDataView,
-    I::Union{OrdinalRange, Colon, AbstractVector{<:Integer}, AbstractVector{<:AbstractString}, Number, AbstractString},
-    J::Union{OrdinalRange, Colon, AbstractVector{<:Integer}, AbstractVector{<:AbstractString}, Number, AbstractString},
-)
-    @boundscheck checkbounds(mu, I, J)
-    i, j = Base.reindex(parentindices(mu), (convertidx(I, mu.obs_names), convertidx(J, mu.var_names)))
-    return view(parent(mu), i, j)
-end
-
-function Base.getindex(
-    mu::MuDataView,
-    I::Union{OrdinalRange, Colon, AbstractVector{<:Integer}, AbstractVector{<:AbstractString}, Number, AbstractString},
-    J::Union{OrdinalRange, Colon, AbstractVector{<:Integer}, AbstractVector{<:AbstractString}, Number, AbstractString},
+    I::Union{OrdinalRange, Colon, AbstractVector{<:Integer}, AbstractVector{<:AbstractString}, Integer, AbstractString},
+    J::Union{OrdinalRange, Colon, AbstractVector{<:Integer}, AbstractVector{<:AbstractString}, Integer, AbstractString},
 )
     @boundscheck checkbounds(mu, I, J)
     i, j = Base.reindex(parentindices(mu), (convertidx(I, mu.obs_names), convertidx(J, mu.var_names)))
-    return getindex(parent(mu), i, j)
+    return @inbounds view(parent(mu), i, j)
 end
 
 Base.parent(mu::MuData) = mu

src/util.jl

@@ -68,17 +68,14 @@ function minimum_unsigned_type_for_n(n::Number)
     return mintype
 end
 
-@inline function convertidx(
-    idx::Union{OrdinalRange, Colon, AbstractVector{<:Integer}},
-    ref::AbstractIndex{<:AbstractString},
-)
+@inline function convertidx(idx::Union{OrdinalRange, Colon, AbstractVector{<:Integer}}, ref::AbstractIndex)
     return idx
 end
-@inline function convertidx(idx::Integer, ref::AbstractIndex{<:AbstractString})
+@inline function convertidx(idx::Integer, ref::AbstractIndex)
     return idx:idx
 end
-@inline convertidx(idx::Union{AbstractString, AbstractVector{<:AbstractString}}, ref::AbstractIndex{<:AbstractString}) =
-    ref[idx, true]
+@inline convertidx(idx::Union{AbstractString, AbstractVector{<:AbstractString}}, ref::AbstractIndex) = ref[idx, true]
+@inline convertidx(idx::AbstractVector{Bool}, ref::AbstractIndex) = findall(idx)
 
 Base.axes(A::Union{AbstractMuData, AbstractAnnData}) = map(n -> Base.OneTo(n), size(A))
 

test/mudata.jl

@@ -127,15 +127,18 @@ function test_row_slice(md, i1, n, d, j=:)
     ad1_names = filter(in(md["ad1"].obs_names), md.obs_names[i1])
     ad2_names = filter(in(md["ad2"].obs_names), md.obs_names[i1])
 
-    ad1_idx = filter(>(0x0), reshape(md.obsmap["ad1"], :)[i1])
-    ad2_idx = filter(>(0x0), reshape(md.obsmap["ad2"], :)[i1])
+    ad1_idx = filter(>(0x0), vec(md.obsmap["ad1"])[i1])
+    ad2_idx = filter(>(0x0), vec(md.obsmap["ad2"])[i1])
+    ad1, ad2 = md["ad1"][ad1_idx, j], md["ad2"][ad2_idx, j]
 
-    ad1, ad2 = parent(md["ad1"])[ad1_idx, j], parent(md["ad2"])[ad2_idx, j]
+    parent_ad1_idx = filter(>(0x0), vec(parent(md).obsmap["ad1"])[parentindices(md)[1][i1]])
+    parent_ad2_idx = filter(>(0x0), vec(parent(md).obsmap["ad2"])[parentindices(md)[1][i1]])
+    parent_ad1, parent_ad2 = parent(md)["ad1"][parent_ad1_idx, j], parent(md["ad2"])[parent_ad2_idx, j]
 
     @test md1["ad1"].X == md2["ad1"].X == md4["ad1"].X == ad1.X
     @test md1["ad2"].X == md2["ad2"].X == md4["ad2"].X == ad2.X
-    @test md1["ad1"].obs_names == md2["ad1"].obs_names == md4["ad1"].obs_names == ad1.obs_names
-    @test md1["ad2"].obs_names == md2["ad2"].obs_names == md4["ad2"].obs_names == ad2.obs_names
+    @test md1["ad1"].obs_names == md2["ad1"].obs_names == md4["ad1"].obs_names == ad1.obs_names == parent_ad1.obs_names
+    @test md1["ad2"].obs_names == md2["ad2"].obs_names == md4["ad2"].obs_names == ad2.obs_names == parent_ad2.obs_names
 end
 
 function test_md_slicing(md, n, d, j=:)
@@ -230,7 +233,7 @@ end
             if unique && axis == 0x2
                 md = (@test_nowarn MuData(mod=Dict("ad1" => ad1, "ad2" => ad2, "ad3" => ad3), axis=axis))
             else
-                md = with_logger(NullLogger()) do    # warning depends on the RNG, and differs between Julia versions
+                md = with_logger(NullLogger()) do # warning depends on the RNG, and differs between Julia versions
                     MuData(mod=Dict("ad1" => ad1, "ad2" => ad2, "ad3" => ad3), axis=axis)
                 end
             end