Add view

aclai-lab · Jul 31, 2024 · 3bb9dbe · 3bb9dbe
1 parent f7a4778
commit 3bb9dbe
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 3 deletions.
diff --git a/src/LabeledMultiDataset.jl b/src/LabeledMultiDataset.jl
@@ -138,7 +138,7 @@ end
 
 function SoleBase.instances(
     lmd::LabeledMultiDataset,
-    inds::AbstractVector{<:Integer},
+    inds::AbstractVector,
     return_view::Union{Val{true},Val{false}} = Val(false)
 )
     LabeledMultiDataset(

diff --git a/src/MultiDataset.jl b/src/MultiDataset.jl
@@ -255,13 +255,23 @@ end
 
 function SoleBase.instances(
     md::MultiDataset,
-    inds::AbstractVector{<:Integer},
+    inds::AbstractVector,
     return_view::Union{Val{true},Val{false}} = Val(false),
 )
     @assert return_view == Val(false)
+    @assert all(i->i<=ninstances(md), inds) "Cannot slice MultiDataset of $(ninstances(md)) instances with indices $(inds)."
     MultiDataset(data(md)[inds,:], grouped_variables(md))
 end
 
+import Base: view
+Base.@propagate_inbounds function view(md::MultiDataset, inds...)
+    MultiDataset(view(data(md), inds...), grouped_variables(md))
+end
+Base.@propagate_inbounds function view(md::MultiDataset, inds::Integer, ::Colon)
+    MultiDataset(view(data(md), [inds], :), grouped_variables(md))
+end
+
+
 function vcat(mds::MultiDataset...)
     MultiDataset(vcat((data.(mds)...)), grouped_variables(first(mds)))
 end

diff --git a/src/dimensional-data.jl b/src/dimensional-data.jl
@@ -47,7 +47,7 @@ function checknvariables(d::AbstractDimensionalDataset{T,D}) where {T<:Number,D}
 end
 nvariables(d::AbstractDimensionalDataset{T,D}) where {T<:Number,D} = size(first(eachinstance(d)), D)
 
-function instances(d::AbstractDimensionalDataset, inds::AbstractVector{<:Integer}, return_view::Union{Val{true},Val{false}} = Val(false))
+function instances(d::AbstractDimensionalDataset, inds::AbstractVector, return_view::Union{Val{true},Val{false}} = Val(false))
     if return_view == Val(true) @views d[inds] else d[inds]    end
 end