Merge pull request #685 from JuliaOpt/slices-of-jumparray

Allow slices of JuMPArray (#684)
jump-dev · Mar 1, 2016 · 50ebf63 · 50ebf63
2 parents 508345e + 06cc476
commit 50ebf63
Show file tree

Hide file tree

Showing 2 changed files with 111 additions and 14 deletions.
diff --git a/src/JuMPArray.jl b/src/JuMPArray.jl
@@ -54,22 +54,38 @@ end
 function _to_cartesian(d,NT,idx...)
     indexing = Any[]
     for (i,S) in enumerate(NT.parameters)
+        idxtype = idx[1][i]
         if S == UnitRange{Int}
-            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)
+            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}
-            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)
+            if idx[1][i] == Colon
+                push!(indexing, Colon())
+            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)
+            end
         else
             push!(indexing, quote
                 if !haskey(d.lookup[$i],idx[$i])

diff --git a/test/variable.jl b/test/variable.jl
@@ -134,3 +134,84 @@ facts("[variable] getValue on empty things") do
     @fact typeof(getValue(z)) --> JuMP.JuMPArray{Float64,3,Tuple{UnitRange{Int},Set{Any},UnitRange{Int}}}
     @fact length(getValue(z)) --> 0
 end
+
+# Slices three-dimensional JuMPContainer x[I+oi,J+oj,K+ok]
+# I,J,K can be singletons, ranges, colons, etc.
+# oi, oj, ok are scalar offsets. The last can be omitted, as
+# that's useful to test more exotic indexing ([:red,"blue"])
+function sliceof(x, I, J, K, oi, oj, ok=nothing)
+    if length(K) == 1
+        if length(J) == 1
+            y = Array(Variable, length(I))
+        else
+            y = Array(Variable, length(I), length(J))
+        end
+    else
+        y = Array(Variable, length(I), length(J), length(K))
+    end
+    ii = 1
+    jj = 1
+    kk = 1
+    for i in I
+        for j in J
+            for k in K
+                if ok == nothing
+                    y[ii,jj,kk] = x[i+oi,j+oj,k]
+                else
+                    y[ii,jj,kk] = x[i+oi,j+oj,k+ok]
+                end
+                kk += 1
+            end
+            jj += 1
+            kk = 1
+        end
+        ii += 1
+        jj = 1
+    end
+    y
+end
+
+facts("[variable] Slices of JuMPArray (#684)") do
+    m = Model()
+    @defVar(m, x[1:3, 1:4,1:2])
+    @defVar(m, y[1:3,-1:2,3:4])
+    @defVar(m, z[1:3,-1:2:4,3:4])
+    @defVar(m, w[1:3,-1:2,[:red,"blue"]])
+
+    @fact x[:] --> vec(sliceof(x, 1:3, 1:4, 1:2, 0, 0, 0))
+    @fact x[:,:,:] --> sliceof(x, 1:3, 1:4, 1:2, 0, 0, 0)
+    @fact x[1,:,:] --> sliceof(x, 1, 1:4, 1:2, 0, 0, 0)
+    @fact x[1,:,2] --> sliceof(x, 1, 1:4, 2, 0, 0, 0)
+    @fact_throws x[1,:,3]
+    @fact x[1:2,:,:] --> sliceof(x, 1:2, 1:4, 1:2, 0, 0, 0)
+    @fact x[1:2,:,2] --> sliceof(x, 1:2, 1:4, 2, 0, 0, 0)
+    @fact x[1:2,:,1:2] --> sliceof(x, 1:2, 1:4, 1:2, 0, 0, 0)
+    @fact_throws x[1:2,:,1:3]
+
+    @fact y[:] --> vec(sliceof(y, 1:3, 1:4, 1:2, 0, -2, 2))
+    @fact y[:,:,:] --> sliceof(y, 1:3, 1:4, 1:2, 0, -2, 2)
+    @fact y[1,:,:] --> sliceof(y, 1, 1:4, 1:2, 0, -2, 2)
+    @fact y[1,:,4] --> sliceof(y, 1, 1:4, 2, 0, -2, 2)
+    @fact_throws y[1,:,5]
+    @fact y[1:2,:,:] --> sliceof(y, 1:2, 1:4, 1:2, 0, -2, 2)
+    @fact y[1:2,:,4] --> sliceof(y, 1:2, 1:4, 2, 0, -2, 2)
+    @fact y[1:2,:,3:4] --> sliceof(y, 1:2, 1:4, 1:2, 0, -2, 2)
+    @fact_throws y[1:2,:,1:3]
+
+    @fact z[:] --> vec(sliceof(z, 1:3, 1:2:6, 1:2, 0, -2, 2))
+    @fact z[:,1,:] --> sliceof(z, 1:3, 3, 1:2, 0, -2, 2)
+    @fact z[1,1,:] --> sliceof(z, 1, 3, 1:2, 0, -2, 2)
+    @fact_throws z[:,5,3]
+    @fact z[1:2,1,:] --> sliceof(z, 1:2, 3, 1:2, 0, -2, 2)
+    @fact z[1:2,1,4] --> sliceof(z, 1:2, 3, 2, 0, -2, 2)
+    @fact z[1:2,1,3:4] --> sliceof(z, 1:2, 3, 1:2, 0, -2, 2)
+    @fact_throws z[1:2,1,1:3]
+
+    @fact w[:] --> vec(sliceof(w, 1:3, 1:4, [:red,"blue"], 0, -2))
+    @fact_throws w[:,:,:]
+    @fact w[1,:,"blue"] --> sliceof(w, 1, 1:4, ["blue"], 0, -2)
+    @fact w[1,:,:red] --> sliceof(w, 1, 1:4, [:red], 0, -2)
+    @fact_throws w[1,:,"green"]
+    @fact w[1:2,:,"blue"] --> sliceof(w, 1:2, 1:4, ["blue"], 0, -2)
+    @fact_throws w[1:2,:,[:red,"blue"]]
+end