Merge pull request #4: Preserve domain information when writing GDX files

Fix: Preserve domain information when writing GDX files

Author: jd-foster

src/GDXFile.jl

@@ -539,6 +539,12 @@ function write_gdx(filepath::String, symbols::Pair{String, DataFrames.DataFrame}
     return filepath
 end
 
+function _set_domain_x(gdx::GDXHandle, name::String, domain::Vector{String}, dim::Int)
+    length(domain) == dim && dim > 0 || return
+    found, sym_nr = gdx_find_symbol(gdx, name)
+    found && gdx_symbol_set_domain_x(gdx, sym_nr, domain)
+end
+
 # Type dispatch for writing symbols
 _write_symbol(gdx::GDXHandle, sym::GDXSet) = _write_set(gdx, sym)
 _write_symbol(gdx::GDXHandle, sym::GDXParameter) = _write_parameter(gdx, sym)
@@ -572,13 +578,14 @@ function _write_set(gdx::GDXHandle, sym::GDXSet)
     end
 
     gdx_data_write_done(gdx)
+    _set_domain_x(gdx, sym.name, sym.domain, dim)
 end
 
 function _write_parameter(gdx::GDXHandle, sym::GDXParameter)
-    _write_parameter_df(gdx, sym.name, sym.records, sym.description)
+    _write_parameter_df(gdx, sym.name, sym.records, sym.description, sym.domain)
 end
 
-function _write_parameter_df(gdx::GDXHandle, name::String, df::DataFrames.DataFrame, description::String="")
+function _write_parameter_df(gdx::GDXHandle, name::String, df::DataFrames.DataFrame, description::String="", domain::Vector{String}=String[])
     dim_cols = [n for n in names(df) if n != "value"]
     dim = length(dim_cols)
 
@@ -596,6 +603,7 @@ function _write_parameter_df(gdx::GDXHandle, name::String, df::DataFrames.DataFr
     end
 
     gdx_data_write_done(gdx)
+    _set_domain_x(gdx, name, domain, dim)
 end
 
 const _VAR_EQU_COLS = Set(["level", "marginal", "lower", "upper", "scale"])
@@ -623,6 +631,7 @@ function _write_variable(gdx::GDXHandle, sym::GDXVariable)
     end
 
     gdx_data_write_done(gdx)
+    _set_domain_x(gdx, sym.name, sym.domain, dim)
 end
 
 function _write_equation(gdx::GDXHandle, sym::GDXEquation)
@@ -648,6 +657,7 @@ function _write_equation(gdx::GDXHandle, sym::GDXEquation)
     end
 
     gdx_data_write_done(gdx)
+    _set_domain_x(gdx, sym.name, sym.domain, dim)
 end
 
 # =============================================================================

src/gdx_c_api.jl

@@ -281,6 +281,22 @@ function gdx_symbol_get_domain_x(gdx::GDXHandle, sym_nr::Int, dim::Int)
     return domains
 end
 
+function gdx_symbol_set_domain_x(gdx_ptr::Ptr{Cvoid}, sym_nr::Int, domain_ids::Vector{String})
+    return ccall(
+        (@cpfx(:gdxsymbolsetdomainx), LIBGDX),
+        Cint,
+        (Ptr{Cvoid}, Cint, Ptr{Cstring}),
+        gdx_ptr,
+        sym_nr,
+        domain_ids,
+    )
+end
+
+function gdx_symbol_set_domain_x(gdx::GDXHandle, sym_nr::Int, domain_ids::Vector{String})
+    rc = gdx_symbol_set_domain_x(gdx.cptr[], sym_nr, domain_ids)
+    return rc
+end
+
 # =============================================================================
 # Alias
 # =============================================================================

test/test_gdxfile.jl

@@ -412,4 +412,114 @@ execute_unload "gams_gdx_test.gdx", i, p, x, y;
         show(io, a)
         @test occursin("j", String(take!(io)))
     end
+
+    @testset "Domain preservation on round-trip (issue #3)" begin
+        set_df = DataFrame(i = ["a", "b", "c"])
+        s = GDXSet("i", "index set", ["*"], set_df)
+        par_df = DataFrame(i = ["a", "b", "c"], value = [10.0, 20.0, 30.0])
+        p = GDXParameter("x", "A parameter over i", ["i"], par_df)
+        gdxfile = GDXFile("", Dict{Symbol,GDXSymbol}(:i => s, :x => p))
+
+        outfile = joinpath(tempdir(), "gdx_jl_domain_test.gdx")
+        write_gdx(outfile, gdxfile)
+
+        gdx2 = read_gdx(outfile)
+        x2 = get_symbol(gdx2, :x)
+        @test x2.domain == ["i"]
+        @test names(gdx2[:x])[1] == "i"
+
+        rm(outfile, force=true)
+    end
+
+    @testset "Domain preservation for variables (issue #3)" begin
+        set_df = DataFrame(i = ["a", "b", "c"])
+        s = GDXSet("i", "index set", ["*"], set_df)
+        var_df = DataFrame(
+            i = ["a", "b", "c"],
+            level = [1.0, 2.0, 3.0],
+            marginal = [0.0, 0.0, 0.0],
+            lower = [-Inf, -Inf, -Inf],
+            upper = [Inf, Inf, Inf],
+            scale = [1.0, 1.0, 1.0]
+        )
+        v = GDXVariable("y", "A variable over i", ["i"], VarFree, var_df)
+        gdxfile = GDXFile("", Dict{Symbol,GDXSymbol}(:i => s, :y => v))
+
+        outfile = joinpath(tempdir(), "gdx_jl_domain_var_test.gdx")
+        write_gdx(outfile, gdxfile)
+
+        gdx2 = read_gdx(outfile)
+        y2 = get_symbol(gdx2, :y)
+        @test y2.domain == ["i"]
+        @test names(gdx2[:y])[1] == "i"
+
+        rm(outfile, force=true)
+    end
+
+    @testset "Domain preservation for equations (issue #3)" begin
+        set_df = DataFrame(i = ["a", "b"])
+        s = GDXSet("i", "index set", ["*"], set_df)
+        eq_df = DataFrame(
+            i = ["a", "b"],
+            level = [1.0, 2.0],
+            marginal = [0.5, 0.6],
+            lower = [-Inf, -Inf],
+            upper = [Inf, Inf],
+            scale = [1.0, 1.0]
+        )
+        eq = GDXEquation("myeq", "test eq", ["i"], EqE, eq_df)
+        gdxfile = GDXFile("", Dict{Symbol,GDXSymbol}(:i => s, :myeq => eq))
+
+        outfile = joinpath(tempdir(), "gdx_jl_domain_eq_test.gdx")
+        write_gdx(outfile, gdxfile)
+
+        gdx2 = read_gdx(outfile)
+        eq2 = get_symbol(gdx2, :myeq)
+        @test eq2.domain == ["i"]
+        @test names(gdx2[:myeq])[1] == "i"
+
+        rm(outfile, force=true)
+    end
+
+    @testset "Multi-dimensional domain preservation (issue #3)" begin
+        si = GDXSet("i", "rows", ["*"], DataFrame(i = ["a", "b"]))
+        sj = GDXSet("j", "cols", ["*"], DataFrame(j = ["x", "y"]))
+        par_df = DataFrame(
+            i = ["a", "a", "b", "b"],
+            j = ["x", "y", "x", "y"],
+            value = [1.0, 2.0, 3.0, 4.0]
+        )
+        p = GDXParameter("cost", "transport cost", ["i", "j"], par_df)
+        gdxfile = GDXFile("", Dict{Symbol,GDXSymbol}(:i => si, :j => sj, :cost => p))
+
+        outfile = joinpath(tempdir(), "gdx_jl_domain_2d_test.gdx")
+        write_gdx(outfile, gdxfile)
+
+        gdx2 = read_gdx(outfile)
+        cost2 = get_symbol(gdx2, :cost)
+        @test cost2.domain == ["i", "j"]
+        @test names(gdx2[:cost])[1:2] == ["i", "j"]
+
+        rm(outfile, force=true)
+    end
+
+    @testset "Domain preservation with GAMS-generated file (issue #3)" begin
+        gdx1 = read_gdx(test_gdx)
+
+        p1 = get_symbol(gdx1, :p)
+        original_domain = p1.domain
+
+        outfile = joinpath(tempdir(), "gdx_jl_gams_domain_rt.gdx")
+        write_gdx(outfile, gdx1)
+
+        gdx2 = read_gdx(outfile)
+        p2 = get_symbol(gdx2, :p)
+        @test p2.domain == original_domain
+
+        x1 = get_symbol(gdx1, :x)
+        x2 = get_symbol(gdx2, :x)
+        @test x2.domain == x1.domain
+
+        rm(outfile, force=true)
+    end
 end