Skip to content

Commit

Permalink
Container revamp (#1099)
Browse files Browse the repository at this point in the history 
Replace JuMPDict with Dict. Rewrite JuMPArray to be compatible with
AbstractArray. Explicit keyword argument in macro to force container
type.
  • Loading branch information
@mlubin
mlubin committed Sep 17, 2017
1 parent 76b3430 commit 497083b
Show file tree
Hide file tree
Showing 10 changed files with 591 additions and 724 deletions.
27 changes: 2 additions & 25 deletions src/JuMP.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,7 @@ export
@objective, @NLobjective,
@NLparameter, @constraintref

include("JuMPContainer.jl")

include("utils.jl")

const MOIVAR = MOI.VariableReference
Expand Down Expand Up @@ -164,7 +164,6 @@ mutable struct Model <: AbstractModel

objdict::Dict{Symbol,Any} # dictionary from variable and constraint names to objects

map_counter::Int # number of times we call getvalue, getdual, getlowerbound and getupperbound on a JuMPContainer, so that we can print out a warning
operator_counter::Int # number of times we add large expressions

# Extension dictionary - e.g. for robust
Expand Down Expand Up @@ -201,7 +200,6 @@ mutable struct Model <: AbstractModel
m.nlpdata = nothing
m.simplify_nonlinear_expressions = simplify_nonlinear_expressions
m.objdict = Dict{Symbol,Any}()
m.map_counter = 0
m.operator_counter = 0
m.ext = Dict{Symbol,Any}()

Expand Down Expand Up @@ -535,14 +533,6 @@ Base.copy(v::Variable, new_model::Model) = Variable(new_model, v.col)
Base.copy(x::Void, new_model::Model) = nothing
Base.copy(v::AbstractArray{Variable}, new_model::Model) = (var -> Variable(new_model, var.col)).(v)

# Copy methods for variable containers
Base.copy(d::JuMPContainer) = map(copy, d)
function Base.copy(d::JuMPContainer, new_model::Model)
new_d = map(x -> copy(x, new_model), d)
new_d.meta[:model] = new_model
new_d
end

##########################################################################
# ConstraintRef
# Reference to a constraint for retrieving solution info
Expand Down Expand Up @@ -760,20 +750,6 @@ function Base.setindex!(m::JuMP.Model, value, name::Symbol)
end

# usage warnings
function mapcontainer_warn(f, x::JuMPContainer, var_or_expr)
isempty(x) && return
v = first(values(x))
m = v.m
m.map_counter += 1
if m.map_counter > 400
# It might not be f that was called the 400 first times but most probably it is f
Base.warn_once("$f has been called on a collection of $(var_or_expr)s a large number of times. For performance reasons, this should be avoided. Instead of $f(x)[a,b,c], use $f(x[a,b,c]) to avoid temporary allocations.")
end
end
mapcontainer_warn(f, x::JuMPContainer{Variable}) = mapcontainer_warn(f, x, "variable")
mapcontainer_warn{E}(f, x::JuMPContainer{E}) = mapcontainer_warn(f, x, "expression")
getvalue_warn(x::JuMPContainer) = nothing

function operator_warn(lhs::AffExpr,rhs::AffExpr)
if length(lhs.vars) > 50 || length(rhs.vars) > 50
if length(lhs.vars) > 1
Expand Down Expand Up @@ -821,6 +797,7 @@ Base.ndims(::JuMPTypes) = 0


##########################################################################
include("containers.jl")
include("operators.jl")
# include("writers.jl")
include("macros.jl")
Expand Down
260 changes: 189 additions & 71 deletions src/JuMPArray.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,96 +3,214 @@
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.

immutable JuMPArray{T,N,NT} <: JuMPContainer{T,N}
innerArray::Array{T,N}
indexsets::NT
lookup::NTuple{N,Any}
meta::Dict{Symbol,Any}
# JuMPArray is inspired by the AxisArrays package.
# JuMPArray can be replaced with AxisArray once integer indices are no longer
# a special case. See discussions at:
# https://github.com/JuliaArrays/AxisArrays.jl/issues/117
# https://github.com/JuliaArrays/AxisArrays.jl/issues/84


struct JuMPArray{T,N,Ax} <: AbstractArray{T,N}
data::Array{T,N}
axes::Ax
lookup::Vector{Dict} # TODO: correctly type return type of the Dict as Int
end

@generated function JuMPArray{T,N}(innerArray::Array{T,N}, indexsets::NTuple{N,Any})
dicttuple = Expr(:tuple)
export JuMPArray

function JuMPArray(data::Array{T,N}, axs...) where {T,N}
lookup = Vector{Dict}(N)
for i in 1:N
inner = quote
idxset = indexsets[$i]
ret = Dict{eltype(idxset), Int}()
end
tupelem = indexsets.parameters[i]
if !(tupelem == UnitRange{Int} || tupelem == StepRange{Int})
inner = quote
$inner
cnt = 1
for x in idxset
ret[x] = cnt
cnt += 1
end
ret
d = Dict{eltype(axs[i]),Int}()
cnt = 1
for el in axs[i]
if haskey(d, el)
error("Repeated index $el. Index sets must have unique elements.")
end
d[el] = cnt
cnt += 1
end
push!(dicttuple.args, inner)
lookup[i] = d
end
:(JuMPArray(innerArray, indexsets, $dicttuple, Dict{Symbol,Any}()))
return JuMPArray(data, axs, lookup)
end

Base.getindex(d::JuMPArray, ::Colon) = d.innerArray[:]
# TODO: use generated function to make this fast
function to_index(A::JuMPArray{T,N}, idx...) where {T,N}
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))...)
end

@generated function Base.getindex{T,N,NT}(d::JuMPArray{T,N,NT}, idx...)
if N != length(idx)
error("Indexed into a JuMPArray with $(length(idx)) indices (expected $N indices)")
# TODO: use generated function to make this fast and type stable
# TODO: better error (or just handle correctly) when user tries to index with a range like a:b
# The only kind of slicing we support is dropping a dimension with colons
function Base.getindex(A::JuMPArray{T}, idx...) where {T}
if Colon() in idx
JuMPArray(A.data[to_index(A,idx...)...], (ax for (i,ax) in enumerate(A.axes) if idx[i] == Colon())...)
else
return A.data[to_index(A,idx...)...]::T
end
Expr(:call, :getindex, :(d.innerArray), _to_cartesian(d,NT,idx)...)
end
Base.getindex(A::JuMPArray, idx::CartesianIndex) = A.data[idx]

Base.setindex!(A::JuMPArray, v, idx...) = A.data[to_index(A,idx...)...] = v
Base.setindex!(A::JuMPArray, v, idx::CartesianIndex) = A.data[idx] = v

# AbstractArray interface

Base.linearindices(A::JuMPArray) = error("linearindices not defined for JuMPArray3")
Base.size(A::JuMPArray) = size(A.data)
Base.indices(A::JuMPArray) = A.axes

@generated function Base.setindex!{T,N,NT}(d::JuMPArray{T,N,NT}, v, idx...)
if N != length(idx)
error("Indexed into a JuMPArray with $(length(idx)) indices (expected $N indices)")
# Arbitrary typed indices. Linear indexing not supported.
struct IndexAnyCartesian <: Base.IndexStyle end
Base.IndexStyle(::Type{JuMPArray{T,N,Ax}}) where {T,N,Ax} = IndexAnyCartesian()

Base.broadcast(f::Function, A::JuMPArray) = JuMPArray(broadcast(f, A.data), A.axes, A.lookup)

Base.isempty(A::JuMPArray) = isempty(A.data)

function Base.isassigned(A::JuMPArray, idx...)
try
to_index(idx...)
return true
catch
return false
end
end
# For ambiguity
function Base.isassigned(A::JuMPArray, idx::Int...)
try
to_index(idx...)
return true
catch
return false
end
Expr(:call, :setindex!, :(d.innerArray), :v, _to_cartesian(d,NT,idx)...)
end

function _to_cartesian(d,NT,idx...)
indexing = Any[]
for (i,S) in enumerate(NT.parameters)
idxtype = idx[1][i]
if S == UnitRange{Int}
if idxtype == Colon
# special stuff
push!(indexing, Colon())
elseif idxtype <: Range
push!(indexing, quote
rng = d.indexsets[$i]
I = idx[$i]
I - (start(rng) - 1)
end)
else
push!(indexing, quote
rng = d.indexsets[$i]
I = idx[$i]
first(rng) <= I <= last(rng) || error("Failed attempt to index JuMPArray along dimension $($i): $I ∉ $(d.indexsets[$i])")
I - (start(rng) - 1)
end)
end
elseif S == StepRange{Int}
if idx[1][i] == Colon
push!(indexing, Colon())
Base.eachindex(A::JuMPArray) = CartesianRange(size(A.data))

# TODO: similar

# Adapted printing from Julia's show.jl

# Copyright (c) 2009-2016: Jeff Bezanson, Stefan Karpinski, Viral B. Shah,
# and other contributors:
#
# https://github.com/JuliaLang/julia/contributors
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.

function summaryio(io::IO, A::JuMPArray)
_summary(io, A)
for (k,ax) in enumerate(A.axes)
print(io, " Dimension $k, ")
show(IOContext(io, :limit=>true), ax)
println(io)
end
print(io, "And data, a ", summary(A.data))
end
_summary(io, A::JuMPArray{T,N}) where {T,N} = println(io, "$N-dimensional JuMPArray{$T,$N,...} with index sets:")

function Base.summary(A::JuMPArray)
io = IOBuffer()
summaryio(io, A)
String(io)
end

function Base.showarray(io::IO, X::JuMPArray, repr::Bool = true; header = true)
repr = false
#if repr && ndims(X) == 1
# return Base.show_vector(io, X, "[", "]")
#end
if !haskey(io, :compact)
io = IOContext(io, :compact => true)
end
if !repr && get(io, :limit, false) && eltype(X) === Method
# override usual show method for Vector{Method}: don't abbreviate long lists
io = IOContext(io, :limit => false)
end
(!repr && header) && print(io, summary(X))
if !isempty(X.data)
(!repr && header) && println(io, ":")
if ndims(X.data) == 0
if isassigned(X.data)
return show(io, X.data[])
else
push!(indexing, quote
rng = $(d.indexsets[i])
I = idx[$i]
first(rng) <= I <= last(rng) || error("Failed attempt to index JuMPArray along dimension $($i): $I ∉ $(d.indexsets[$i])")
dv, rv = divrem(I - start(rng), step(rng))
rv == 0 || error("Failed attempt to index JuMPArray along dimension $($i): $I ∉ $(d.indexsets[$i])")
dv + 1
end)
return print(io, undef_ref_str)
end
end
#if repr
# if ndims(X.data) <= 2
# Base.print_matrix_repr(io, X)
# else
# show_nd(io, X, print_matrix_repr, false)
# end
#else
punct = (" ", " ", "")
if ndims(X.data) <= 2
Base.print_matrix(io, X.data, punct...)
else
push!(indexing, quote
if !haskey(d.lookup[$i],idx[$i])
error("Failed attempt to index JuMPArray along dimension $($i): $(idx[$i]) ∉ $(d.indexsets[$i])")
show_nd(io, X,
(io, slice) -> Base.print_matrix(io, slice, punct...),
!repr)
end
#end
elseif repr
Base.repremptyarray(io, X.data)
end
end

# n-dimensional arrays
function show_nd(io::IO, a::JuMPArray, print_matrix, label_slices)
limit::Bool = get(io, :limit, false)
if isempty(a)
return
end
tailinds = Base.tail(Base.tail(indices(a.data)))
nd = ndims(a)-2
for I in CartesianRange(tailinds)
idxs = I.I
if limit
for i = 1:nd
ii = idxs[i]
ind = tailinds[i]
if length(ind) > 10
if ii == ind[4] && all(d->idxs[d]==first(tailinds[d]),1:i-1)
for j=i+1:nd
szj = size(a,j+2)
indj = tailinds[j]
if szj>10 && first(indj)+2 < idxs[j] <= last(indj)-3
@goto skip
end
end
#println(io, idxs)
print(io, "...\n\n")
@goto skip
end
if ind[3] < ii <= ind[end-3]
@goto skip
end
end
d.lookup[$i][idx[$i]]::Int
end)
end
end
if label_slices
print(io, "[:, :, ")
for i = 1:(nd-1); show(io, a.axes[i+2][idxs[i]]); print(io,", "); end
show(io, a.axes[end][idxs[end]])
println(io, "] =")
end
slice = view(a.data, indices(a.data,1), indices(a.data,2), idxs...)
Base.print_matrix(io, slice)
print(io, idxs == map(last,tailinds) ? "" : "\n\n")
@label skip
end
indexing
end
Loading

0 comments on commit 497083b

Please sign in to comment.