Remove object references from model.ext[:miplearn]

src/instance/jump.jl

@@ -23,26 +23,22 @@ function __init_PyJuMPInstance__()
 
         function get_variable_features(self, var_name)
             model = self.model
-            v = variable_by_name(model, var_name)
+            return get(model.ext[:miplearn][:variable_features], var_name, nothing)
-            return get(model.ext[:miplearn][:variable_features], v, nothing)
         end
 
         function get_variable_category(self, var_name)
             model = self.model
-            v = variable_by_name(model, var_name)
+            return get(model.ext[:miplearn][:variable_categories], var_name, nothing)
-            return get(model.ext[:miplearn][:variable_categories], v, nothing)
         end
 
         function get_constraint_features(self, cname)
             model = self.model
-            c = constraint_by_name(model, cname)
+            return get(model.ext[:miplearn][:constraint_features], cname, nothing)
-            return get(model.ext[:miplearn][:constraint_features], c, nothing)
         end
 
         function get_constraint_category(self, cname)
             model = self.model
-            c = constraint_by_name(model, cname)
+            return get(model.ext[:miplearn][:constraint_categories], cname, nothing)
-            return get(model.ext[:miplearn][:constraint_categories], c, nothing)
         end
     end
     copy!(PyJuMPInstance, Class)
@@ -78,23 +74,12 @@ function save(filename::AbstractString, instance::JuMPInstance)::Nothing
         miplearn.write_pickle_gz(instance.py.samples, py_samples_filename, quiet=true)
         py_samples = read(py_samples_filename)
 
-        # Replace variable/constraint refs by names
-        _to_names(d) = Dict(name(var) => value for (var, value) in d)
-        ext_original = instance.model.ext[:miplearn]
-        ext_names = Dict(
-            :variable_features => _to_names(ext_original[:variable_features]),
-            :variable_categories => _to_names(ext_original[:variable_categories]),
-            :constraint_features => _to_names(ext_original[:constraint_features]),
-            :constraint_categories => _to_names(ext_original[:constraint_categories]),
-            :instance_features => ext_original[:instance_features],
-        )
-
         # Generate JLD2 file
         jldsave(
             filename;
             miplearn_version="0.2",
             mps=mps,
-            ext=ext_names,
+            ext=instance.model.ext[:miplearn],
             py_samples=py_samples,
         )
     end
@@ -124,30 +109,13 @@ function load_instance(filename::AbstractString)::JuMPInstance
             mps_filename = "$(tempname()).mps.gz"
             write(mps_filename, file["mps"])
             model = read_from_file(mps_filename)
+            model.ext[:miplearn] = file["ext"]
 
             # Unpickle instance.py.samples
             py_samples_filename = tempname()
             write(py_samples_filename, file["py_samples"])
             py_samples = miplearn.read_pickle_gz(py_samples_filename, quiet=true)
 
-            # Replace variable/constraint names by refs
-            _to_var(model, d) = Dict(
-                variable_by_name(model, varname) => value
-                for (varname, value) in d
-            )
-            _to_constr(model, d) = Dict(
-                constraint_by_name(model, cname) => value
-                for (cname, value) in d
-            )
-            ext = file["ext"]
-            model.ext[:miplearn] = Dict(
-                :variable_features => _to_var(model, ext[:variable_features]),
-                :variable_categories => _to_var(model, ext[:variable_categories]),
-                :constraint_features => _to_constr(model, ext[:constraint_features]),
-                :constraint_categories => _to_constr(model, ext[:constraint_categories]),
-                :instance_features => ext[:instance_features],
-            )
-
             instance = JuMPInstance(model)
             instance.py.samples = py_samples
         end

src/solvers/macros.jl

@@ -6,10 +6,10 @@ function init_miplearn_ext(model)::Dict
     if :miplearn ∉ keys(model.ext)
         model.ext[:miplearn] = Dict{Symbol, Any}()
         model.ext[:miplearn][:instance_features] = [0.0]
-        model.ext[:miplearn][:variable_features] = Dict{VariableRef, Vector{Float64}}()
+        model.ext[:miplearn][:variable_features] = Dict{AbstractString, Vector{Float64}}()
-        model.ext[:miplearn][:variable_categories] = Dict{VariableRef, String}()
+        model.ext[:miplearn][:variable_categories] = Dict{AbstractString, String}()
-        model.ext[:miplearn][:constraint_features] = Dict{ConstraintRef, Vector{Float64}}()
+        model.ext[:miplearn][:constraint_features] = Dict{AbstractString, Vector{Float64}}()
-        model.ext[:miplearn][:constraint_categories] = Dict{ConstraintRef, String}()
+        model.ext[:miplearn][:constraint_categories] = Dict{AbstractString, String}()
     end
     return model.ext[:miplearn]
 end
@@ -24,28 +24,32 @@ end
 
 function set_features!(v::VariableRef, f::Array{Float64})::Nothing
     ext = init_miplearn_ext(v.model)
-    ext[:variable_features][v] = f
+    n = _get_and_check_name(v)
+    ext[:variable_features][n] = f
     return
 end
 
 
 function set_category!(v::VariableRef, category::String)::Nothing
     ext = init_miplearn_ext(v.model)
-    ext[:variable_categories][v] = category
+    n = _get_and_check_name(v)
+    ext[:variable_categories][n] = category
     return
 end
 
 
 function set_features!(c::ConstraintRef, f::Array{Float64})::Nothing
     ext = init_miplearn_ext(c.model)
-    ext[:constraint_features][c] = f
+    n = _get_and_check_name(c)
+    ext[:constraint_features][n] = f
     return
 end
 
 
 function set_category!(c::ConstraintRef, category::String)::Nothing
     ext = init_miplearn_ext(c.model)
-    ext[:constraint_categories][c] = category
+    n = _get_and_check_name(c)
+    ext[:constraint_categories][n] = category
     return
 end
 
@@ -63,6 +67,15 @@ macro category(obj, category)
     end
 end
 
+function _get_and_check_name(obj)
+    n = name(obj)
+    length(n) > 0 || error(
+        "Features and categories can only be assigned to variables and " *
+        "constraints that have names. Unnamed model element detected."
+    )
+    return n
+end
+
 
 export @feature,
        @category

test/fixtures/knapsack.jl

@@ -28,14 +28,14 @@ function build_knapsack_model()
     end
 
     # Should store ML information
-    @test model.ext[:miplearn][:variable_features][x[1]] == [1.0, 5.0]
+    @test model.ext[:miplearn][:variable_features]["x[1]"] == [1.0, 5.0]
-    @test model.ext[:miplearn][:variable_features][x[2]] == [2.0, 6.0]
+    @test model.ext[:miplearn][:variable_features]["x[2]"] == [2.0, 6.0]
-    @test model.ext[:miplearn][:variable_features][x[3]] == [3.0, 7.0]
+    @test model.ext[:miplearn][:variable_features]["x[3]"] == [3.0, 7.0]
-    @test model.ext[:miplearn][:variable_categories][x[1]] == "type-1"
+    @test model.ext[:miplearn][:variable_categories]["x[1]"] == "type-1"
-    @test model.ext[:miplearn][:variable_categories][x[2]] == "type-2"
+    @test model.ext[:miplearn][:variable_categories]["x[2]"] == "type-2"
-    @test model.ext[:miplearn][:variable_categories][x[3]] == "type-3"
+    @test model.ext[:miplearn][:variable_categories]["x[3]"] == "type-3"
-    @test model.ext[:miplearn][:constraint_features][c1] == [1.0, 2.0, 3.0]
+    @test model.ext[:miplearn][:constraint_features]["c1"] == [1.0, 2.0, 3.0]
-    @test model.ext[:miplearn][:constraint_categories][c1] == "c1"
+    @test model.ext[:miplearn][:constraint_categories]["c1"] == "c1"
     @test model.ext[:miplearn][:instance_features] == [5.0]    
 
     return model