close #68

jump-dev · Feb 22, 2015 · 92c93fd · 92c93fd
1 parent e54a8ce
commit 92c93fd
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 1 deletion.
diff --git a/src/atoms/affine/index.jl b/src/atoms/affine/index.jl
@@ -1,4 +1,4 @@
-import Base.getindex
+import Base: getindex, to_index
 export IndexAtom, getindex
 
 ArrayOrNothing = Union(AbstractArray, Nothing)
@@ -80,3 +80,9 @@ getindex(x::AbstractExpr, ind::Real) = getindex(x, ind:ind)
 getindex(x::AbstractExpr, row::Real, col::Real) = getindex(x, row:row, col:col)
 getindex{T <: Real}(x::AbstractExpr, row::Real, cols::AbstractArray{T, 1}) = getindex(x, row:row, cols)
 getindex{T <: Real}(x::AbstractExpr, rows::AbstractArray{T, 1}, col::Real) = getindex(x, rows, col:col)
+function getindex(x::AbstractExpr, I::AbstractArray{Bool,2})
+    return [ x[i] for i in to_index(I) ]
+end
+function getindex(x::AbstractExpr, I::AbstractVector{Bool})
+    return [ x[i] for i in to_index(I) ]
+end
diff --git a/test/test_affine.jl b/test/test_affine.jl
@@ -122,6 +122,14 @@ facts("Affine Atoms") do
     @fact p.optval => roughly(2, TOL)
     @fact evaluate(x[1] + x[2])[1] => roughly(2, TOL)
 
+    x = Variable(3)
+    I = [true true false]
+    p = minimize(sum(x[I]), [x >= 1])
+    @fact vexity(p) => AffineVexity()
+    solve!(p)
+    @fact p.optval => roughly(2, TOL)
+    @fact evaluate(sum(x[I]))[1] => roughly(2, TOL)
+
     rows = 6
     cols = 8
     n = 2