diff --git a/src/ADNIDatasets.jl b/src/ADNIDatasets.jl index bd3c987..7d7ed7d 100644 --- a/src/ADNIDatasets.jl +++ b/src/ADNIDatasets.jl @@ -15,14 +15,14 @@ struct ADNIScanData end struct ADNISubject - ID::Int64 - n_scans::Int64 + ID::Int + n_scans::Int scan_dates::Vector{Date} Data::Vector{ADNIScanData} end struct ADNIDataset - n_subjects::Int64 + n_subjects::Int SubjectData::Vector{ADNISubject} rois::Vector{String} end @@ -209,6 +209,11 @@ function Base.getindex(data::ADNIDataset, idx::Vector{Int}) return ADNIDataset(length(idx), data.SubjectData[idx], data.rois) end +function Base.getindex(data::ADNIDataset, idx::UnitRange{Int}) + _idx = collect(idx) + return ADNIDataset(length(idx), data.SubjectData[_idx], data.rois) +end + function Base.iterate(d::ADNIDataset, state=1) state > length(d) ? nothing : (d[state], state+1) end @@ -221,6 +226,8 @@ function Base.length(data::ADNIDataset) get(data, @lens _.n_subjects) end +Base.lastindex(d::ADNIDataset) = length(d) + function Base.filter(func, data::ADNIDataset) d = Iterators.filter(func, data) |> collect ADNIDataset(length(d), d, data.rois) @@ -237,6 +244,11 @@ function Base.getindex(sub::ADNISubject, idx::Vector{Int}) return ADNISubject(sub.ID, length(idx), sub.scan_dates[idx], sub.Data[idx]) end +function Base.getindex(sub::ADNISubject, idx::UnitRange{Int}) + _idx = collect(idx) + return ADNISubject(sub.ID, length(idx), sub.scan_dates[_idx], sub.Data[_idx]) +end + function Base.iterate(d::ADNISubject, state=1) state > length(d) ? nothing : (d[state], state+1) end @@ -245,6 +257,8 @@ function Base.length(data::ADNISubject) get(data, @lens _.n_scans) end +Base.lastindex(d::ADNISubject) = length(d) + # Exports export ADNIDataset, ADNISubject, ADNIScanData export get_suvr, get_ref_suvr, get_ref_vol, get_vol, get_dates, get_times, diff --git a/test/runtests.jl b/test/runtests.jl index c099b66..ab37afc 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -49,10 +49,12 @@ data_scans_2 = ADNIDataset(datadf; min_scans=2, max_scans=2, qc=false) @test _idx_subset isa ADNIDataset @test get_id.(_idx_subset) == [2, 3] @test get_id.(_idx_subset) == get_id.(_filter_subset) + @test get_id.(data_scans_1[2:end]) == get_id.(data_scans_1[[2,3]]) # testing indexing of ADNISubject @test data_scans_1[2][1] isa ADNIScanData @test data_scans_1[2][[1,2]] isa ADNISubject @test length(data_scans_1[2][[1,2]]) == 2 @test data_scans_1[2][[1,2]][1].Date == data_scans_1[2][1].Date + @test allequal(get_dates(data_scans_1[2][[2,3]]) .== get_dates(data_scans_1[2][2:end])) end \ No newline at end of file