Support indexing by arbitrary sets

README.md

@@ -40,24 +40,25 @@ platform is supported.
 # Defining variables
 
 Variables are defined using the ``@defVar`` macro. The first argument must be a ``Model``.
+In the examples below we assume ``m`` is already defined.
 The second is an expression that declares the variable name and optinally allows specification
 of lower and upper bounds. The possible combinations are
 
 No bounds:
 
-	x
+	@defVar(m, x )
 
 Lower bound only:
 
-	x >= lb
+	@defVar(m, x >= lb )
 
 Upper bound only:
 
-	x <= ub
+	@defVar(m, x <= ub )
 
 Lower and upper bounds:
 
-	lb <= x <= ub
+	@defVar(m, lb <= x <= ub )
 
 Where ``x`` must be a valid symbol which will be assigned to in the local context.
 
@@ -66,9 +67,15 @@ Integer and binary restrictions can optionally be specified with a third argumen
 Arrays of variable objects can created by appending brackets to the variable name.
 For example,
 
-	x[1:M,1:N] >= 0
+	@defVar(m, x[1:M,1:N] >= 0 )
 
-Will create an ``M`` by ``N`` array of variables. Currently only ranges starting at ``1`` are supported as index sets, but this will be extended in the future.
+will create an ``M`` by ``N`` array of variables. Both ranges and arbitrary iterable sets are supported as index sets. Using ranges will generally be faster, as this avoids dictionary lookups. The following code
+
+	s = ["Green","Blue"]
+	@defVar(m, x[-10:10,s] , Int)
+	x[-4,"Green"]
+
+works fine.
 
 Bounds can depend on variable indices:
 

src/MathProg.jl

@@ -5,6 +5,10 @@
 # By Iain Dunning and Miles Lubin
 ########################################################################
 
+import Base.getindex
+import Base.setindex!
+
+
 module MathProg
 
 # Currently we use CLP and CoinMP, but later we will want to make it so

src/macros.jl

@@ -247,9 +247,9 @@ macro defVar(m, x, extra...)
         idxvars = Symbol[]
         idxsets = {}
         arrcall = Expr(:call,:Array,Variable)
-        refcall = Expr(:ref,varname)
+        refcall = Expr(:ref,esc(varname))
         for s in var.args[2:end]
-            if s.head == :(=)
+            if isa(s,Expr) && s.head == :(=)
                 idxvar = s.args[1]
                 idxset = s.args[2]
             else
@@ -258,25 +258,27 @@ macro defVar(m, x, extra...)
             end
             push!(idxvars, idxvar)
             push!(idxsets, idxset)
-            push!(arrcall.args, :(length($(idxset))))
-            push!(refcall.args, idxvar) 
-            if !isa(idxset,Expr) || idxset.head != :(:) || idxset.args[1] != 1
-                error("Expected $idxset to be a range starting at 1. (Syntax will be expanded in the future)")
-            end
+            push!(arrcall.args, :(length($(esc(idxset)))))
+            push!(refcall.args, esc(idxvar)) 
         end
-        code = :( $(refcall) = MathProg.Variable($m, $lb, $ub, $t) )
+        code = :( $(refcall) = MathProg.Variable($(esc(m)), $(esc(lb)), $(esc(ub)), $t) )
         for (idxvar, idxset) in zip(reverse(idxvars),reverse(idxsets))
             code = quote
-                for $(idxvar) in $idxset
-                    $code
+                let
+                    for $(esc(idxvar)) in $(esc(idxset))
+                        $code
+                    end
                 end
             end
         end
+
+        mac = Expr(:macrocall,symbol("@gendict"),varname,:Variable,idxsets...)
         code = quote 
-            $(varname) = MultivarDict($arrcall,$(string(varname)))
+            #$(esc(varname)) = MultivarDict($arrcall,$(string(varname)))
+            $(esc(mac))
             $code
         end
-        return esc(code)
+        return code
     end
 end
 

src/multivardict.jl

@@ -5,5 +5,87 @@ type MultivarDict{T,N}
 end
 
 getindex(d::MultivarDict, vals...) = getindex(d.innerArray, vals...)
-
 setindex!(d::MultivarDict, vals...) = setindex!(d.innerArray, vals...)
+
+
+# generate and instantiate a type which is indexed by the given index sets
+# the following types of index sets are allowed:
+# 0:K -- range with compile-time starting index
+# S -- general iterable set
+macro gendict(instancename,T,idxsets...)
+    N = length(idxsets)
+    typename = symbol(string("MathProgDict",gensym()))
+    isrange = Array(Bool,N)
+    offset = Array(Int,N)
+    dictnames = Array(Symbol,N)
+    for i in 1:N
+        if isa(idxsets[i],Expr) && idxsets[i].head == :(:)
+            isrange[i] = true
+            if isa(idxsets[i].args[1],Int)
+                offset[i] = 1 - idxsets[i].args[1]
+            else
+                error("Currently only ranges with integer compile-time starting values are allowed as index sets. $(idxsets[i].args[1]) is not an integer in range $(idxsets[i]).")
+            end
+        else
+            isrange[i] = false
+            dictnames[i] = gensym()
+        end
+    end
+    #typecode = :(type $(esc(typename)); innerArray::Array{$T,$N}; name::String; end)
+    typecode = :(type $(typename); innerArray::Array{$T,$N}; name::String; end)
+    builddicts = quote end
+    for i in 1:N
+        if !isrange[i]
+            push!(typecode.args[3].args,:($(dictnames[i])::Dict))
+            push!(builddicts.args, quote 
+                $(esc(dictnames[i])) = Dict(); 
+                let 
+                    for (j,k) in enumerate($(esc(idxsets[i])))
+                        $(esc(dictnames[i]))[k] = j
+                    end 
+                end end)
+        end
+    end
+    getidxlhs = :(getindex(d::$(typename)))
+    setidxlhs = :(setindex!(d::$(typename),val::$T))
+    getidxrhs = :(getindex(d.innerArray))
+    setidxrhs = :(setindex!(d.innerArray,val))
+    for i in 1:N
+        varname = symbol(string("x",i))
+
+        if isrange[i]
+            push!(getidxlhs.args,:($varname::Int))
+            push!(setidxlhs.args,:($varname::Int))
+
+            push!(getidxrhs.args,:($varname+$(offset[i])))
+            push!(setidxrhs.args,:($varname+$(offset[i])))
+        else
+            push!(getidxlhs.args,varname)
+            push!(setidxlhs.args,varname)
+
+            push!(getidxrhs.args,:(d.($(Expr(:quote,dictnames[i])))[$varname]))
+            push!(setidxrhs.args,:(d.($(Expr(:quote,dictnames[i])))[$varname]))
+        end
+    end
+
+    #funcs = :($(esc(getidxlhs)) = $(esc(getidxrhs)); $(esc(setidxlhs)) = $(esc(setidxrhs)))
+    funcs = :($getidxlhs = $getidxrhs; $setidxlhs = $setidxrhs)
+    geninstance = :($(esc(instancename)) = $(typename)(Array($T),$(string(instancename))))
+    for i in 1:N
+        push!(geninstance.args[2].args[2].args, :(length($(esc(idxsets[i])))))
+        if !isrange[i]
+            push!(geninstance.args[2].args, esc(dictnames[i]))
+        end
+    end
+    eval(Expr(:toplevel, typecode))
+    eval(Expr(:toplevel, funcs))
+
+    quote
+        $builddicts
+        $geninstance
+    end
+
+end
+
+export @gendict
+