Implement save, load_jump_instance

src/modeling/jump_instance.jl

@@ -3,6 +3,7 @@
 #  Released under the modified BSD license. See COPYING.md for more details.
 
 using JuMP
+using JLD2
 
 
 @pydef mutable struct PyJuMPInstance <: miplearn.Instance
@@ -23,38 +24,149 @@ using JuMP
     function get_variable_features(self, var_name)
         model = self.model
         v = variable_by_name(model, var_name)
-        return get(model.ext[:miplearn][:variable_features], v, [0.0])
+        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], v, "default")
+        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], c, [0.0])
+        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], c, "default")
+        return get(model.ext[:miplearn][:constraint_categories], c, nothing)
     end
 end
 
 
 struct JuMPInstance
     py::PyCall.PyObject
+    model::Model
 end
 
 
 function JuMPInstance(model)
     model isa Model || error("model should be a JuMP.Model. Found $(typeof(model)) instead.")
-    return JuMPInstance(PyJuMPInstance(model))
+    return JuMPInstance(
+        PyJuMPInstance(model),
+        model,
+    )
 end
 
 
-export JuMPInstance
+function save(filename::AbstractString, instance::JuMPInstance)::Nothing
+    # Convert JuMP model to MPS
+    mps_filename = "$(tempname()).mps.gz"
+    write_to_file(instance.model, mps_filename)
+    mps = read(mps_filename)
+
+    # Pickle instance.py.samples. Ideally, we would use dumps and loads, but this
+    # causes some issues with PyCall, probably due to automatic type conversions.
+    py_samples_filename = tempname()
+    miplearn.write_pickle_gz(instance.py.samples, py_samples_filename)
+    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,
+        py_samples=py_samples,
+    )
+    return
+end
+
+
+function load_jump_instance(filename::AbstractString)::JuMPInstance
+    jldopen(filename, "r") do file
+        file["miplearn_version"] == 0.2 || error(
+            "MIPLearn version 0.2 cannot read instance files generated by " *
+            "version $(file["miplearn_version"])."
+        )
+
+        # Convert MPS to JuMP
+        mps_filename = "$(tempname()).mps.gz"
+        write(mps_filename, file["mps"])
+        model = read_from_file(mps_filename)
+
+        # Unpickle instance.py.samples
+        py_samples_filename = tempname()
+        write(py_samples_filename, file["py_samples"])
+        py_samples = miplearn.read_pickle_gz(py_samples_filename)
+
+        # 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
+
+        return instance
+    end
+end
+
+
+struct FileInstance
+    filename::AbstractString
+    loaded::Union{Nothing,JuMPInstance}
+end
+
+
+function FileInstance(filename::AbstractString)::FileInstance
+    return FileInstance(
+        filename,
+        nothing,
+    )
+end
+
+
+function load!(instance::FileInstance)
+    instance.loaded = load_jump_instance(instance.filename)
+end
+
+
+function free!(instance::FileInstance)
+    instance.loaded = nothing
+end
+
+
+function flush!(instance::FileInstance)
+    save(instance.filename, instance.loaded)
+end
+
+
+export JuMPInstance, save, load_jump_instance