Wed, 15 Dec 2021 01:09:09 +0200
Implement ZipArrays
36 | 1 | ####################################### |
2 | # Simultaneous view to multiple arrays | |
3 | ####################################### | |
4 | ||
5 | """ | |
6 | Implements a type for zipped arrays; see `ZipArray`, `AbstractZipArray`, | |
7 | `ZipVector`, and `AbstractZipVector`. | |
8 | """ | |
9 | module ZipArrays | |
10 | ||
11 | export ZipArray, | |
12 | ZipVector, | |
13 | AbstractZipArray, | |
14 | AbstractZipVector | |
15 | ||
16 | """ | |
17 | `ZipArray{S,T,N,A,B}` | |
18 | ||
19 | `N`-dimensional array of element type `Tuple{S,T}`, the component of the tuple stored in separate | |
20 | arrays of type `A <: AbstractArray{S,N}` and `B <: AbstractArray{T,N}`. | |
21 | """ | |
22 | struct ZipArray{S,T,N,A <: AbstractArray{S,N}, B <: AbstractArray{T,N}} <: AbstractArray{Tuple{S,T}, N} | |
23 | a :: A | |
24 | b :: B | |
25 | ||
26 | function ZipArray(a :: A, b :: B) where {S,T,N,A <: AbstractArray{S,N}, B <: AbstractArray{T,N}} | |
27 | @assert(IndexStyle(A)==IndexStyle(B)) | |
28 | @assert(size(a)==size(b)) | |
29 | return new{S,T,N,A,B}(a, b) | |
30 | end | |
31 | ||
32 | end | |
33 | ||
34 | """ | |
35 | `AbstractZipArray{S,T,N}` | |
36 | ||
37 | A way to refer to a `ZipArray` without specifying the container types. | |
38 | """ | |
39 | AbstractZipArray{S,T,N} = ZipArray{S,T,N,A,B} where {A <: AbstractArray{S,N}, B <: AbstractArray{T,N}} | |
40 | ||
41 | """ | |
42 | `ZipVector{S,T,A,B}` | |
43 | ||
44 | One-dimensional `ZipArray`. | |
45 | """ | |
46 | ZipVector{S,T,A,B} = ZipArray{S,T,1,A,B} | |
47 | ||
48 | """ | |
49 | `AbstractZipVector{S,T}` | |
50 | ||
51 | One-dimensional `AbstractZipArray`. | |
52 | """ | |
53 | AbstractZipVector{S,T} = AbstractZipArray{S,T,1} | |
54 | ||
55 | ||
56 | function Base.IndexStyle( :: Type{<:ZipArray{S,T,N,A,B}}) where {S,T,N,A,B} | |
57 | return IndexStyle(A) | |
58 | end | |
59 | ||
60 | @inline function Base.size(z :: AbstractZipArray{S,T,N}) where {S,T,N} | |
61 | return size(z.a) | |
62 | end | |
63 | ||
64 | @inline function Base.getindex(z :: AbstractZipArray{S,T,N}, i) where {S,T,N} | |
65 | return (getindex(z.a, i), getindex(z.b, i)) | |
66 | end | |
67 | ||
68 | @inline function Base.setindex!(z :: AbstractZipArray{S,T,N}, (s, t) :: Tuple{S,T}, i) where {S,T,N} | |
69 | return (setindex!(z.a, s, i), setindex!(z.b, t, i)) | |
70 | end | |
71 | ||
72 | @inline function Base.pop!(z :: AbstractZipVector{S,T}) where {S,T} | |
73 | s = pop!(z.a) | |
74 | t = pop!(z.b) | |
75 | return (s, t) | |
76 | end | |
77 | ||
78 | @inline function Base.push!(z :: AbstractZipVector{S,T}, items :: Vararg{Tuple{S,T},M}) where {S,T,M} | |
79 | push!(z.a, s for (s, _) ∈ items) | |
80 | push!(z.b, t for (_, t) ∈ items) | |
81 | end | |
82 | ||
83 | end # module | |
84 |