|
3 | 3 | # License, v. 2.0. If a copy of the MPL was not distributed with this |
4 | 4 | # file, You can obtain one at http://mozilla.org/MPL/2.0/. |
5 | 5 |
|
6 | | -immutable JuMPArray{T,N,NT} <: JuMPContainer{T,N} |
7 | | - innerArray::Array{T,N} |
8 | | - indexsets::NT |
9 | | - lookup::NTuple{N,Any} |
10 | | - meta::Dict{Symbol,Any} |
| 6 | +# JuMPArray is inspired by the AxisArrays package. |
| 7 | +# JuMPArray can be replaced with AxisArray once integer indices are no longer |
| 8 | +# a special case. See discussions at: |
| 9 | +# https://github.com/JuliaArrays/AxisArrays.jl/issues/117 |
| 10 | +# https://github.com/JuliaArrays/AxisArrays.jl/issues/84 |
| 11 | + |
| 12 | + |
| 13 | +struct JuMPArray{T,N,Ax} <: AbstractArray{T,N} |
| 14 | + data::Array{T,N} |
| 15 | + axes::Ax |
| 16 | + lookup::Vector{Dict} # TODO: correctly type return type of the Dict as Int |
11 | 17 | end |
12 | 18 |
|
13 | | -@generated function JuMPArray{T,N}(innerArray::Array{T,N}, indexsets::NTuple{N,Any}) |
14 | | - dicttuple = Expr(:tuple) |
| 19 | +export JuMPArray |
| 20 | + |
| 21 | +function JuMPArray(data::Array{T,N}, axs...) where {T,N} |
| 22 | + lookup = Vector{Dict}(N) |
15 | 23 | for i in 1:N |
16 | | - inner = quote |
17 | | - idxset = indexsets[$i] |
18 | | - ret = Dict{eltype(idxset), Int}() |
19 | | - end |
20 | | - tupelem = indexsets.parameters[i] |
21 | | - if !(tupelem == UnitRange{Int} || tupelem == StepRange{Int}) |
22 | | - inner = quote |
23 | | - $inner |
24 | | - cnt = 1 |
25 | | - for x in idxset |
26 | | - ret[x] = cnt |
27 | | - cnt += 1 |
28 | | - end |
29 | | - ret |
| 24 | + d = Dict{eltype(axs[i]),Int}() |
| 25 | + cnt = 1 |
| 26 | + for el in axs[i] |
| 27 | + if haskey(d, el) |
| 28 | + error("Repeated index $el. Index sets must have unique elements.") |
30 | 29 | end |
| 30 | + d[el] = cnt |
| 31 | + cnt += 1 |
31 | 32 | end |
32 | | - push!(dicttuple.args, inner) |
| 33 | + lookup[i] = d |
33 | 34 | end |
34 | | - :(JuMPArray(innerArray, indexsets, $dicttuple, Dict{Symbol,Any}())) |
| 35 | + return JuMPArray(data, axs, lookup) |
35 | 36 | end |
36 | 37 |
|
37 | | -Base.getindex(d::JuMPArray, ::Colon) = d.innerArray[:] |
| 38 | +# TODO: use generated function to make this fast |
| 39 | +function to_index(A::JuMPArray{T,N}, idx...) where {T,N} |
| 40 | + return tuple((isa(i,Colon) ? Colon() : (k <= N ? A.lookup[k][i] : (((i == 1) ? 1 : error("invalid index $i")))) for (k,i) in enumerate(idx))...) |
| 41 | +end |
38 | 42 |
|
39 | | -@generated function Base.getindex{T,N,NT}(d::JuMPArray{T,N,NT}, idx...) |
40 | | - if N != length(idx) |
41 | | - error("Indexed into a JuMPArray with $(length(idx)) indices (expected $N indices)") |
| 43 | +# TODO: use generated function to make this fast and type stable |
| 44 | +# TODO: better error (or just handle correctly) when user tries to index with a range like a:b |
| 45 | +# The only kind of slicing we support is dropping a dimension with colons |
| 46 | +function Base.getindex(A::JuMPArray{T}, idx...) where {T} |
| 47 | + if Colon() in idx |
| 48 | + JuMPArray(A.data[to_index(A,idx...)...], (ax for (i,ax) in enumerate(A.axes) if idx[i] == Colon())...) |
| 49 | + else |
| 50 | + return A.data[to_index(A,idx...)...]::T |
42 | 51 | end |
43 | | - Expr(:call, :getindex, :(d.innerArray), _to_cartesian(d,NT,idx)...) |
44 | 52 | end |
| 53 | +Base.getindex(A::JuMPArray, idx::CartesianIndex) = A.data[idx] |
| 54 | + |
| 55 | +Base.setindex!(A::JuMPArray, v, idx...) = A.data[to_index(A,idx...)...] = v |
| 56 | +Base.setindex!(A::JuMPArray, v, idx::CartesianIndex) = A.data[idx] = v |
| 57 | + |
| 58 | +# AbstractArray interface |
| 59 | + |
| 60 | +Base.linearindices(A::JuMPArray) = error("JuMPArray does not support this operation.") |
| 61 | +Base.size(A::JuMPArray) = error("JuMPArray does not define this operation") |
| 62 | +Base.indices(A::JuMPArray) = A.axes |
45 | 63 |
|
46 | | -@generated function Base.setindex!{T,N,NT}(d::JuMPArray{T,N,NT}, v, idx...) |
47 | | - if N != length(idx) |
48 | | - error("Indexed into a JuMPArray with $(length(idx)) indices (expected $N indices)") |
| 64 | +# Arbitrary typed indices. Linear indexing not supported. |
| 65 | +struct IndexAnyCartesian <: Base.IndexStyle end |
| 66 | +Base.IndexStyle(::Type{JuMPArray{T,N,Ax}}) where {T,N,Ax} = IndexAnyCartesian() |
| 67 | + |
| 68 | +Base.broadcast(f::Function, A::JuMPArray) = JuMPArray(broadcast(f, A.data), A.axes, A.lookup) |
| 69 | + |
| 70 | +Base.isempty(A::JuMPArray) = isempty(A.data) |
| 71 | + |
| 72 | +function Base.isassigned(A::JuMPArray, idx...) |
| 73 | + try |
| 74 | + to_index(idx...) |
| 75 | + return true |
| 76 | + catch |
| 77 | + return false |
| 78 | + end |
| 79 | +end |
| 80 | +# For ambiguity |
| 81 | +function Base.isassigned(A::JuMPArray, idx::Int...) |
| 82 | + try |
| 83 | + to_index(idx...) |
| 84 | + return true |
| 85 | + catch |
| 86 | + return false |
49 | 87 | end |
50 | | - Expr(:call, :setindex!, :(d.innerArray), :v, _to_cartesian(d,NT,idx)...) |
51 | 88 | end |
52 | 89 |
|
53 | | -function _to_cartesian(d,NT,idx...) |
54 | | - indexing = Any[] |
55 | | - for (i,S) in enumerate(NT.parameters) |
56 | | - idxtype = idx[1][i] |
57 | | - if S == UnitRange{Int} |
58 | | - if idxtype == Colon |
59 | | - # special stuff |
60 | | - push!(indexing, Colon()) |
61 | | - elseif idxtype <: Range |
62 | | - push!(indexing, quote |
63 | | - rng = d.indexsets[$i] |
64 | | - I = idx[$i] |
65 | | - I - (start(rng) - 1) |
66 | | - end) |
67 | | - else |
68 | | - push!(indexing, quote |
69 | | - rng = d.indexsets[$i] |
70 | | - I = idx[$i] |
71 | | - first(rng) <= I <= last(rng) || error("Failed attempt to index JuMPArray along dimension $($i): $I ∉ $(d.indexsets[$i])") |
72 | | - I - (start(rng) - 1) |
73 | | - end) |
74 | | - end |
75 | | - elseif S == StepRange{Int} |
76 | | - if idx[1][i] == Colon |
77 | | - push!(indexing, Colon()) |
| 90 | +Base.eachindex(A::JuMPArray) = CartesianRange(size(A.data)) |
| 91 | + |
| 92 | +# TODO: similar |
| 93 | + |
| 94 | +# Adapted printing from Julia's show.jl |
| 95 | + |
| 96 | +# Copyright (c) 2009-2016: Jeff Bezanson, Stefan Karpinski, Viral B. Shah, |
| 97 | +# and other contributors: |
| 98 | +# |
| 99 | +# https://github.com/JuliaLang/julia/contributors |
| 100 | +# |
| 101 | +# Permission is hereby granted, free of charge, to any person obtaining |
| 102 | +# a copy of this software and associated documentation files (the |
| 103 | +# "Software"), to deal in the Software without restriction, including |
| 104 | +# without limitation the rights to use, copy, modify, merge, publish, |
| 105 | +# distribute, sublicense, and/or sell copies of the Software, and to |
| 106 | +# permit persons to whom the Software is furnished to do so, subject to |
| 107 | +# the following conditions: |
| 108 | +# |
| 109 | +# The above copyright notice and this permission notice shall be |
| 110 | +# included in all copies or substantial portions of the Software. |
| 111 | + |
| 112 | +function summaryio(io::IO, A::JuMPArray) |
| 113 | + _summary(io, A) |
| 114 | + for (k,ax) in enumerate(A.axes) |
| 115 | + print(io, " Dimension $k, ") |
| 116 | + show(IOContext(io, :limit=>true), ax) |
| 117 | + println(io) |
| 118 | + end |
| 119 | + print(io, "And data, a ", summary(A.data)) |
| 120 | +end |
| 121 | +_summary(io, A::JuMPArray{T,N}) where {T,N} = println(io, "$N-dimensional JuMPArray{$T,$N,...} with index sets:") |
| 122 | + |
| 123 | +function Base.summary(A::JuMPArray) |
| 124 | + io = IOBuffer() |
| 125 | + summaryio(io, A) |
| 126 | + String(io) |
| 127 | +end |
| 128 | + |
| 129 | +function Base.showarray(io::IO, X::JuMPArray, repr::Bool = true; header = true) |
| 130 | + repr = false |
| 131 | + #if repr && ndims(X) == 1 |
| 132 | + # return Base.show_vector(io, X, "[", "]") |
| 133 | + #end |
| 134 | + if !haskey(io, :compact) |
| 135 | + io = IOContext(io, :compact => true) |
| 136 | + end |
| 137 | + if !repr && get(io, :limit, false) && eltype(X) === Method |
| 138 | + # override usual show method for Vector{Method}: don't abbreviate long lists |
| 139 | + io = IOContext(io, :limit => false) |
| 140 | + end |
| 141 | + (!repr && header) && print(io, summary(X)) |
| 142 | + if !isempty(X.data) |
| 143 | + (!repr && header) && println(io, ":") |
| 144 | + if ndims(X.data) == 0 |
| 145 | + if isassigned(X.data) |
| 146 | + return show(io, X.data[]) |
78 | 147 | else |
79 | | - push!(indexing, quote |
80 | | - rng = $(d.indexsets[i]) |
81 | | - I = idx[$i] |
82 | | - first(rng) <= I <= last(rng) || error("Failed attempt to index JuMPArray along dimension $($i): $I ∉ $(d.indexsets[$i])") |
83 | | - dv, rv = divrem(I - start(rng), step(rng)) |
84 | | - rv == 0 || error("Failed attempt to index JuMPArray along dimension $($i): $I ∉ $(d.indexsets[$i])") |
85 | | - dv + 1 |
86 | | - end) |
| 148 | + return print(io, undef_ref_str) |
87 | 149 | end |
| 150 | + end |
| 151 | + #if repr |
| 152 | + # if ndims(X.data) <= 2 |
| 153 | + # Base.print_matrix_repr(io, X) |
| 154 | + # else |
| 155 | + # show_nd(io, X, print_matrix_repr, false) |
| 156 | + # end |
| 157 | + #else |
| 158 | + punct = (" ", " ", "") |
| 159 | + if ndims(X.data) <= 2 |
| 160 | + Base.print_matrix(io, X.data, punct...) |
88 | 161 | else |
89 | | - push!(indexing, quote |
90 | | - if !haskey(d.lookup[$i],idx[$i]) |
91 | | - error("Failed attempt to index JuMPArray along dimension $($i): $(idx[$i]) ∉ $(d.indexsets[$i])") |
| 162 | + show_nd(io, X, |
| 163 | + (io, slice) -> Base.print_matrix(io, slice, punct...), |
| 164 | + !repr) |
| 165 | + end |
| 166 | + #end |
| 167 | + elseif repr |
| 168 | + Base.repremptyarray(io, X.data) |
| 169 | + end |
| 170 | +end |
| 171 | + |
| 172 | +# n-dimensional arrays |
| 173 | +function show_nd(io::IO, a::JuMPArray, print_matrix, label_slices) |
| 174 | + limit::Bool = get(io, :limit, false) |
| 175 | + if isempty(a) |
| 176 | + return |
| 177 | + end |
| 178 | + tailinds = Base.tail(Base.tail(indices(a.data))) |
| 179 | + nd = ndims(a)-2 |
| 180 | + for I in CartesianRange(tailinds) |
| 181 | + idxs = I.I |
| 182 | + if limit |
| 183 | + for i = 1:nd |
| 184 | + ii = idxs[i] |
| 185 | + ind = tailinds[i] |
| 186 | + if length(ind) > 10 |
| 187 | + if ii == ind[4] && all(d->idxs[d]==first(tailinds[d]),1:i-1) |
| 188 | + for j=i+1:nd |
| 189 | + szj = size(a,j+2) |
| 190 | + indj = tailinds[j] |
| 191 | + if szj>10 && first(indj)+2 < idxs[j] <= last(indj)-3 |
| 192 | + @goto skip |
| 193 | + end |
| 194 | + end |
| 195 | + #println(io, idxs) |
| 196 | + print(io, "...\n\n") |
| 197 | + @goto skip |
| 198 | + end |
| 199 | + if ind[3] < ii <= ind[end-3] |
| 200 | + @goto skip |
| 201 | + end |
92 | 202 | end |
93 | | - d.lookup[$i][idx[$i]]::Int |
94 | | - end) |
| 203 | + end |
| 204 | + end |
| 205 | + if label_slices |
| 206 | + print(io, "[:, :, ") |
| 207 | + for i = 1:(nd-1); show(io, a.axes[i+2][idxs[i]]); print(io,", "); end |
| 208 | + show(io, a.axes[end][idxs[end]]) |
| 209 | + println(io, "] =") |
95 | 210 | end |
| 211 | + slice = view(a.data, indices(a.data,1), indices(a.data,2), idxs...) |
| 212 | + Base.print_matrix(io, slice) |
| 213 | + print(io, idxs == map(last,tailinds) ? "" : "\n\n") |
| 214 | + @label skip |
96 | 215 | end |
97 | | - indexing |
98 | 216 | end |
0 commit comments