Add HDF5 group to codec dispatch

JuliaQUBO · Dec 4, 2023 · 968a857 · 968a857
1 parent 4838337
commit 968a857
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 25 deletions.
diff --git a/src/library/format/qubin/parser.jl b/src/library/format/qubin/parser.jl
@@ -4,7 +4,7 @@ function read_model(path::AbstractString, fmt::QUBin)
     end
 end
 
-function read_model(fp::HDF5.File, fmt::QUBin)
+function read_model(fp::P, fmt::QUBin) where {P<:Union{HDF5.File,HDF5.Group}}
     form      = _parse_model_form(fp, fmt)
     variables = _parse_model_variables(fp, fmt)
     metadata  = _parse_model_metadata(fp, fmt)
@@ -19,7 +19,7 @@ function read_solution(path::AbstractString, fmt::QUBin)
     end
 end
 
-function read_solution(fp::HDF5.File, fmt::QUBin)
+function read_solution(fp::P, fmt::QUBin) where {P<:Union{HDF5.File,HDF5.Group}}
     return _parse_solution(fp, fmt)
 end
 
@@ -32,7 +32,7 @@ function _parse_model(
     return Model{V,T,U}(variable_map, form; metadata, solution)
 end
 
-function _parse_model_form(fp::HDF5.File, ::QUBin)
+function _parse_model_form(fp::P, ::QUBin) where {P<:Union{HDF5.File,HDF5.Group}}
     n = read(fp["model"]["form"]["dimension"])
 
     li = read(fp["model"]["form"]["linear"]["i"])
@@ -67,7 +67,7 @@ function _parse_model_form(
     return Form{T}(n, L, Q, α, β; sense, domain)
 end
 
-function _parse_model_variables(fp::HDF5.File, ::QUBin)
+function _parse_model_variables(fp::P, ::QUBin) where {P<:Union{HDF5.File,HDF5.Group}}
     variables = read(fp["model"]["variables"])
 
     return _parse_model_variables(variables)
@@ -77,11 +77,11 @@ function _parse_model_variables(variables::Vector{V}) where {V}
     return VariableMap{V}(variables)
 end
 
-function _parse_model_metadata(fp::HDF5.File, ::QUBin)
+function _parse_model_metadata(fp::P, ::QUBin) where {P<:Union{HDF5.File,HDF5.Group}}
     return JSON.parse(read(fp["model"]["metadata"]))
 end
 
-function _parse_solution(fp::HDF5.File, fmt::QUBin)
+function _parse_solution(fp::P, fmt::QUBin) where {P<:Union{HDF5.File,HDF5.Group}}
     data     = _parse_solution_data(fp, fmt)
     metadata = _parse_solution_metadata(fp, fmt)
 
@@ -100,7 +100,7 @@ function _parse_solution(
     return SampleSet{T,U}(data; metadata, sense, domain)
 end
 
-function _parse_solution_data(fp::HDF5.File, ::QUBin)
+function _parse_solution_data(fp::P, ::QUBin) where {P<:Union{HDF5.File,HDF5.Group}}
     ψ = read(fp["solution"]["data"]["state"])
     λ = read(fp["solution"]["data"]["value"])
     r = read(fp["solution"]["data"]["reads"])
@@ -112,6 +112,6 @@ function _parse_solution_data(ψ::Matrix{U}, λ::Vector{T}, r::Vector{Int}) wher
     return Sample{T,U}[Sample{T,U}(ψ[i, :], λ[i], r[i]) for i in eachindex(λ)]
 end
 
-function _parse_solution_metadata(fp::HDF5.File, ::QUBin)
+function _parse_solution_metadata(fp::P, ::QUBin) where {P<:Union{HDF5.File,HDF5.Group}}
     return JSON.parse(read(fp["solution"]["metadata"]))
 end
diff --git a/src/library/format/qubin/printer.jl b/src/library/format/qubin/printer.jl
@@ -11,17 +11,21 @@ function write_model(
 end
 
 function write_model(
-    fp::HDF5.File,
+    fp::P,
     model::M,
     fmt::QUBin,
-) where {V,T,U,M<:AbstractModel{V,T,U}}
+) where {P<:Union{HDF5.File,HDF5.Group},V,T,U,M<:AbstractModel{V,T,U}}
     _write_model(fp, model, fmt)
     _write_solution(fp, model, fmt)
 
     return nothing
 end
 
-function write_solution(path::AbstractString, sol::S, fmt::QUBin) where {S<:AbstractSolution}
+function write_solution(
+    path::AbstractString,
+    sol::S,
+    fmt::QUBin,
+) where {S<:AbstractSolution}
     HDF5.h5open(path, "w") do fp
         write_solution(fp, sol, fmt)
     end
@@ -36,10 +40,10 @@ function write_solution(fp::HDF5.File, sol::S, fmt::QUBin) where {S<:AbstractSol
 end
 
 function _write_model(
-    fp::HDF5.File,
+    fp::P,
     model::M,
     fmt::QUBin,
-) where {V,T,U,M<:AbstractModel{V,T,U}}
+) where {P<:Union{HDF5.File,HDF5.Group},V,T,U,M<:AbstractModel{V,T,U}}
     HDF5.create_group(fp, "model")
 
     _write_model_variables(fp, model, fmt)
@@ -50,20 +54,20 @@ function _write_model(
 end
 
 function _write_model_variables(
-    fp::HDF5.File,
+    fp::P,
     model::M,
     ::QUBin,
-) where {V,T,U,M<:AbstractModel{V,T,U}}
+) where {P<:Union{HDF5.File,HDF5.Group},V,T,U,M<:AbstractModel{V,T,U}}
     fp["model"]["variables"] = variables(model)
 
     return nothing
 end
 
 function _write_model_form(
-    fp::HDF5.File,
+    fp::P,
     model::M,
     ::QUBin,
-) where {V,T,U,M<:AbstractModel{V,T,U}}
+) where {P<:Union{HDF5.File,HDF5.Group},V,T,U,M<:AbstractModel{V,T,U}}
     HDF5.create_group(fp["model"], "form")
 
     n, L, Q, α, β, s, x = QUBOTools.form(model, :sparse)
@@ -96,20 +100,20 @@ end
 
 
 function _write_model_metadata(
-    fp::HDF5.File,
+    fp::P,
     model::M,
     ::QUBin,
-) where {V,T,U,M<:AbstractModel{V,T,U}}
+) where {P<:Union{HDF5.File,HDF5.Group},V,T,U,M<:AbstractModel{V,T,U}}
     fp["model"]["metadata"] = JSON.json(QUBOTools.metadata(model))
 
     return nothing
 end
 
 function _write_solution(
-    fp::HDF5.File,
+    fp::P,
     model::M,
     fmt::QUBin,
-) where {V,T,U,M<:AbstractModel{V,T,U}}
+) where {P<:Union{HDF5.File,HDF5.Group},V,T,U,M<:AbstractModel{V,T,U}}
     HDF5.create_group(fp, "solution")
 
     sol = QUBOTools.solution(model)
@@ -125,10 +129,10 @@ function _write_solution(
 end
 
 function _write_solution_data(
-    fp::HDF5.File,
+    fp::P,
     sol::AbstractSolution{T,U},
     ::QUBin,
-) where {T,U}
+) where {P<:Union{HDF5.File,HDF5.Group},T,U}
     HDF5.create_group(fp["solution"], "data")
 
     if isempty(sol)
@@ -149,10 +153,10 @@ function _write_solution_data(
 end
 
 function _write_solution_metadata(
-    fp::HDF5.File,
+    fp::P,
     sol::AbstractSolution{T,U},
     ::QUBin,
-) where {T,U}
+) where {P<:Union{HDF5.File,HDF5.Group},T,U}
     fp["solution"]["metadata"] = JSON.json(QUBOTools.metadata(sol))
 
     return nothing