optimizations for selecting a row or col from a matrix exp

src/atoms/index.jl

@@ -6,58 +6,85 @@ getindex{T <: Real}(x::Constant, inds::AbstractArray{T, 1}) = Constant(getindex(
 getindex{T <: Real}(x::Constant, rows::AbstractArray{T, 1}, cols::AbstractArray{T, 1}) = Constant(getindex(x.value, rows, cols))
 
 function getindex{T <: Real}(x::AffineExpr, inds::AbstractArray{T, 1})
-  # number of rows/cols in the coefficient for x in our canonical form
-  num_rows_coeff = length(inds)
-  num_cols_coeff = x.size[1] * x.size[2]
-  indexer = Constant(sparse(1:length(inds), inds, 1.0, num_rows_coeff, num_cols_coeff))
+  if length(inds) == 1
+    # select only one entry of x
+    ind = convert(Int64, inds[1])
+    vars_to_coeffs_map = Dict{Uint64, Constant}()
+    for (v, c) in x.vars_to_coeffs_map
+      vars_to_coeffs_map[v] = c[ind, :]
+    end
+    constant = x.constant[ind, :]
+  else
+    # number of rows/cols in the coefficient for x in our canonical form
+    num_rows_coeff = length(inds)
+    num_cols_coeff = x.size[1] * x.size[2]
+    indexer = Constant(sparse(1:length(inds), inds, 1.0, num_rows_coeff, num_cols_coeff))
 
-  vars_to_coeffs_map = Dict{Uint64, Constant}()
-  for (v, c) in x.vars_to_coeffs_map
-    vars_to_coeffs_map[v] = indexer * c
+    vars_to_coeffs_map = Dict{Uint64, Constant}()
+    for (v, c) in x.vars_to_coeffs_map
+      vars_to_coeffs_map[v] = indexer * c
+    end
+    constant = indexer * x.constant
   end
-  constant = indexer * x.constant
-  this = AffineExpr(:index, (x,), vars_to_coeffs_map, constant, (num_rows_coeff, 1))
+  this = AffineExpr(:index, (x,), vars_to_coeffs_map, constant, (length(inds), 1))
   this.evaluate = ()->x.evaluate()[inds]
   return this
 end
 
-function getindex{T <: Real}(x::AffineExpr, row::T, col::T)
-  position = x.size[1] * (convert(Int64, col) - 1) + convert(Int64, row)
-  vars_to_coeffs_map = Dict{Uint64, Constant}()
-  for (v, c) in x.vars_to_coeffs_map
-    vars_to_coeffs_map[v] = c[position, :]
-  end
-  constant = x.constant[position, :]
-  this = AffineExpr(:index, (x,), vars_to_coeffs_map, constant, (1, 1))
-  this.evaluate = ()->x.evaluate()[row, col]
-  return this
-end
-
 function getindex{T <: Real}(x::AffineExpr, rows::AbstractArray{T, 1}, cols::AbstractArray{T, 1})
   if length(rows) == 1 && length(cols) == 1
-    return getindex(x, rows[1], cols[1])
-  end
-  # number of rows/cols in the coefficient for x in our canonical form
-  num_rows_coeff = length(rows) * length(cols)
-  num_cols_coeff = x.size[1] * x.size[2]
+    # select only one entry of x
+    ind = x.size[1] * (convert(Int64, cols[1]) - 1) + convert(Int64, rows[1])
+    vars_to_coeffs_map = Dict{Uint64, Constant}()
+    for (v, c) in x.vars_to_coeffs_map
+      vars_to_coeffs_map[v] = c[ind, :]
+    end
+    constant = x.constant[ind, :]
+  elseif length(cols) == 1 && typeof(rows) == UnitRange{Int64}
+    # build a range object to just select the part of a column in x
+    ind = x.size[1] * (convert(Int64, cols[1]) - 1)
+    ind_range = (ind + rows[1]) : (ind + rows[end])
+    vars_to_coeffs_map = Dict{Uint64, Constant}()
 
-  # Create the indexing matrix such that indexer * vec(x) = vec(x[rows, cols])
-  J = Array(Int64, num_rows_coeff)
-  k = 1
-  num_rows = x.size[1]
-  for c in cols
-    for r in rows
-      J[k] = num_rows * (convert(Int64, c) - 1) + convert(Int64, r)
-      k += 1
+    for (v, c) in x.vars_to_coeffs_map
+      vars_to_coeffs_map[v] = c[ind_range, :]
     end
-  end
-  indexer = Constant(sparse(1:num_rows_coeff, J, 1.0, num_rows_coeff, num_cols_coeff))
+    constant = x.constant[ind_range, :]
+  elseif length(rows) == 1 && typeof(cols) == UnitRange{Int64}
+    # build a range object to just select part of a row in x
+    start_ind = x.size[1] * (cols[1] - 1)
+    end_ind = x.size[1] * (cols[end] - 1)
+    ind_range = (start_ind + rows[1]) : x.size[1] : (end_ind + rows[1])
+
+    vars_to_coeffs_map = Dict{Uint64, Constant}()
+    for (v, c) in x.vars_to_coeffs_map
+      vars_to_coeffs_map[v] = c[ind_range, :]
+    end
+    constant = x.constant[ind_range, :]
+  else
+    # number of rows/cols in the coefficient for x in our canonical form
+    num_rows_coeff = length(rows) * length(cols)
+    num_cols_coeff = x.size[1] * x.size[2]
 
-  vars_to_coeffs_map = Dict{Uint64, Constant}()
-  for (v, c) in x.vars_to_coeffs_map
-    vars_to_coeffs_map[v] = indexer * c
+    # Create the indexing matrix such that indexer * vec(x) = vec(x[rows, cols])
+    J = Array(Int64, num_rows_coeff)
+    k = 1
+    num_rows = x.size[1]
+    for c in cols
+      for r in rows
+        J[k] = num_rows * (convert(Int64, c) - 1) + convert(Int64, r)
+        k += 1
+      end
+    end
+    indexer = Constant(sparse(1:num_rows_coeff, J, 1.0, num_rows_coeff, num_cols_coeff))
+
+    vars_to_coeffs_map = Dict{Uint64, Constant}()
+    for (v, c) in x.vars_to_coeffs_map
+      vars_to_coeffs_map[v] = indexer * c
+    end
+    constant = indexer * x.constant
   end
-  constant = indexer * x.constant
+
   this = AffineExpr(:index, (x,), vars_to_coeffs_map, constant, (length(rows), length(cols)))
   this.evaluate = ()->x.evaluate()[rows, cols]
   return this

test/unit_tests/index.jl

@@ -4,8 +4,14 @@ y = x[1:3, 5:end]
 z = x[:]
 t = [1, 2, 3]
 f = t * x[2, 3]
+row = x[1, :]
+col = x[:, 1]
+entry = x[4]
 test_expr(y, (x,))
 test_expr(z, (x,))
 test_expr(f, (x,))
+test_expr(row, (x,))
+test_expr(col, (x,))
+test_expr(entry, (x,))
 
 info("All index tests passed")