Merge pull request #175 from psrenergy/px/negative-states-alt

Fix support for negative states + Add tests

src/OpenBinary/writer.jl

@@ -3,6 +3,9 @@ Base.@kwdef mutable struct Writer <: PSRI.AbstractWriter
 
     # stages
     stage_total::Int
+    first_stage::Int
+    last_stage::Int
+    num_stages::Int
     # first_year::Int
     initial_stage::Int
     # first_relative_stage::Int
@@ -57,6 +60,8 @@ function PSRI.open(
     stages::Integer = 0,
     blocks::Integer = 0,
     scenarios::Integer = 0,
+    first_stage::Integer = 1,
+    last_stage::Integer = first_stage + stages - 1,
     agents::Vector{String} = String[],
     unit::Union{Nothing, String} = nothing,
     # optional
@@ -159,6 +164,7 @@ function PSRI.open(
     end
 
     directory = dirname(path)
+
     if !isdir(directory)
         error("Directory $directory does not exist.")
     end
@@ -185,13 +191,14 @@ function PSRI.open(
         version = 4
     end
 
+    num_stages = last_stage - first_stage + 1
+
     write(ioh, Int32(0))
     write(ioh, Int32(version))
     write(ioh, Int32(0))
     write(ioh, Int32(0))
-    first_relative_stage = 1 # TODO
-    write(ioh, Int32(first_relative_stage)) # relative stage
-    write(ioh, Int32(stages))
+    write(ioh, Int32(first_stage))
+    write(ioh, Int32(last_stage))
     write(ioh, Int32(scenarios))
     write(ioh, Int32(length(agents)))
     write(ioh, Int32(scenarios_type))
@@ -231,7 +238,8 @@ function PSRI.open(
         write(ioh, Int32(0))
         write(ioh, Int32(0)) # offset1
         write(ioh, Int32(blocks)) # offset2 -> block_total = offset2 - offset1
-        for i in first_relative_stage:stages-1
+
+        for _ in first_stage:(num_stages-1)
             write(ioh, Int32(0))
         end
     else
@@ -241,7 +249,7 @@ function PSRI.open(
         write(ioh, Int32(0))
         acc = Int32(0)
         blocks = 0
-        for t in first_relative_stage:stages
+        for t in first_stage:stages
             # TODO write hourly files
             b = PSRI.blocks_in_stage(
                 is_hourly,
@@ -290,6 +298,9 @@ function PSRI.open(
     end
 
     return Writer(;
+        first_stage = first_stage,
+        last_stage = last_stage,
+        num_stages = num_stages,
         io = io,
         offset = header_size,
         stage_total = stages,
@@ -309,6 +320,10 @@ function PSRI.open(
     )
 end
 
+function _get_stage_offset(ior::Writer)
+    return 4 * ior.agents_total * ior.scenario_total * (ior.first_stage - 1)
+end
+
 PSRI.is_hourly(graf::Writer) = graf.is_hourly
 PSRI.hour_discretization(graf::Writer) = graf.hour_discretization
 PSRI.stage_type(graf::Writer) = graf.stage_type
@@ -328,21 +343,24 @@ function PSRI.write_registry(
         error("File is not in open state.")
     end
 
-    if !(1 <= stage <= io.stage_total)
-        error("stage should be between 1 and $(io.stage_total)")
+    if !(io.first_stage <= stage <= io.last_stage)
+        error("stage should be between '$(io.first_stage)' and '$(io.last_stage)'")
     end
 
     if !(1 <= scenario <= io.scenario_total)
-        error("scenarios should be between 1 and $(io.scenario_total)")
+        error("scenarios should be between '1' and '$(io.scenario_total)'")
     end
 
     blocks_in_stage = PSRI.blocks_in_stage(io, stage)
+
     if !(1 <= block <= blocks_in_stage) # io.blocks
-        error("block should be between 1 and $blocks_in_stage")
+        error("block should be between '1' and '$blocks_in_stage'")
     end
 
     if length(data) != io.agents_total
-        error("data vector has length $(length(data)) and expected was $(io.agents_total)")
+        error(
+            "data vector has length '$(length(data))' and expected was '$(io.agents_total)'",
+        )
     end
 
     current = position(io.io)
@@ -366,38 +384,53 @@ function _reopen_pre_write(io::Writer)
         io.io = Base.open(io.FILE_PATH, "a")
         io.is_open = true
     end
+
     return nothing
 end
+
 function _reopen_pos_write(io::Writer)
     if io.reopen_mode
         Base.close(io.io)
         io.is_open = false
     end
+
     return nothing
 end
 
+function _get_registry_size(io)
+    return sizeof(Float32) * io.agents_total
+end
+
+function _get_last_position(io)
+    return _get_position(
+        io,
+        io.last_stage,
+        io.scenario_total,
+        io.is_hourly ? io.blocks_per_stage[end] : io.block_total,
+    ) + _get_registry_size(io)
+end
+
 function PSRI.close(io::Writer)
     io.is_open = false
     io.reopen_mode = false # so that it wont try to reopen
+
     seekend(io.io)
-    current = position(io.io)
-    last =
-        _get_position(
-            io,
-            io.stage_total,
-            io.scenario_total,
-            io.is_hourly ? io.blocks_per_stage[end] : io.block_total,
-        ) + 4 * io.agents_total
-
-    if current != last
-        seek(io.io, last - 4 * io.agents_total)
+
+    curr_pos = position(io.io)
+    last_pos = _get_last_position(io)
+
+    if curr_pos != last_pos
+        seek(io.io, last_pos - _get_registry_size(io))
+
         write(io.io, Float32(0))
-        println(
-            "File not writen completely. Expected $(div(last, 4)) registries, got $(div(current, 4))",
+
+        @warn(
+            "File not writen completely. Expected $(div(last, 4)) registries, got $(div(current, 4))"
         )
     end
 
     Base.close(io.io)
+
     return nothing
 end
 

src/OpenStudy/duration.jl

@@ -145,7 +145,7 @@ function _variable_duration_to_file!(data::Data)
         OpenBinary.Reader,
         FILE_NAME;
         use_header = false,
-        first_stage = data.first_date,
+        initial_stage = data.first_date,
     )
 
     data.variable_duration = ior
@@ -242,13 +242,13 @@ function _hour_block_map_to_file!(data::Data)
         OpenBinary.Reader,
         FILE_NAME_DUR;
         use_header = false,
-        first_stage = data.first_date,
+        initial_stage = data.first_date,
     )
     ior_hbm = open(
         OpenBinary.Reader,
         FILE_NAME_HBM;
         use_header = false,
-        first_stage = data.first_date,
+        initial_stage = data.first_date,
     )
 
     data.variable_duration = ior_dur

src/reader_mapper.jl

@@ -41,7 +41,7 @@ function add_reader!(
         T,
         path;
         header = header,
-        first_stage = mapper.first_date,
+        initial_stage = mapper.first_date,
     )
 
     push!(mapper.list, graf)

src/tables/interface.jl

@@ -99,17 +99,17 @@ struct GrafTable{T} <: Tables.AbstractColumns
 
             i = 0
 
-            for t in 1:reader.stage_total, s in 1:reader.scenario_total
+            for t in reader.first_stage:reader.last_stage, s in 1:reader.scenario_total
                 for b in 1:reader.blocks_per_stage[t]
                     i += 1
 
+                    goto(reader, t, s, b)
+
                     domain[i, :] .= [t, s, b]
 
                     for a in 1:reader.agents_total
                         matrix[i, a] = reader[a]
                     end
-
-                    next_registry(reader)
                 end
             end
         else
@@ -119,19 +119,19 @@ struct GrafTable{T} <: Tables.AbstractColumns
 
             i = 0
 
-            for t in 1:reader.stage_total,
+            for t in reader.first_stage:reader.last_stage,
                 s in 1:reader.scenario_total,
                 b in 1:reader.block_total
 
                 i += 1
 
+                goto(reader, t, s, b)
+
                 domain[i, :] .= [t, s, b]
 
                 for a in 1:reader.agents_total
                     matrix[i, a] = reader[a]
                 end
-
-                next_registry(reader)
             end
         end
 

test/OpenBinary/nonpositive_indices.jl

@@ -1,23 +1,67 @@
 function test_nonpositive_indices()
-    case_path = joinpath(@__DIR__, "..", "data", "case5", "inflow")
+    path = joinpath(@__DIR__, "..", "data", "case5", "inflow")
 
-    ior = PSRI.open(
-        PSRI.OpenBinary.Reader,
-        case_path;
-        use_header = false,
-    )
+    io_r = PSRI.open(PSRI.OpenBinary.Reader, path; use_header = false)
 
-    @test ior isa PSRI.OpenBinary.Reader
-    @test ior.first_stage == 1
-    @test ior.first_relative_stage == -5
-    @test ior.relative_stage_skip == 0
+    @test io_r isa PSRI.OpenBinary.Reader
+    @test io_r.initial_stage == 5
+    @test io_r.first_stage == -2
+    @test io_r.relative_stage_skip == 0
     @test PSRI.OpenBinary._get_position(
-        ior,
-        ior.first_relative_stage,
+        io_r,
+        io_r.first_stage,
         1,
         1,
     ) == 0
 
+    src_table = CSV.read(
+        "$path.csv", DataFrames.DataFrame;
+        header = 4,
+        skipto = 5,
+    )
+
+    dst_table = DataFrames.DataFrame(
+        PSRI.GrafTable{Float64}(path; use_header = false),
+    )
+
+    @testset "Read" begin
+        @test PSRI.goto(io_r, -2, 1, 1) === nothing
+
+        @test_throws AssertionError PSRI.goto(io_r, -3, 1, 1)
+
+        @test size(src_table) == size(dst_table) # Check size first
+
+        @test Matrix(src_table) ≈ Matrix(dst_table)
+    end
+
+    temp_path = tempname()
+
+    io_w = PSRI.open(
+        PSRI.OpenBinary.Writer,
+        temp_path;
+        first_stage = -2,
+        unit = io_r.unit,
+        stages = io_r.stage_total,
+        blocks = io_r.block_total,
+        scenarios = io_r.scenario_total,
+        agents = io_r.agent_names,
+    )
+
+    @testset "Write" begin
+        @test io_w.first_stage == -2
+
+        for row in eachrow(src_table)
+            t, s, b, data... = row
+
+            cache = collect(Float64, data)
+
+            PSRI.write_registry(io_w, cache, t, s, b)
+        end
+    end
+
+    PSRI.close(io_w)
+    PSRI.close(io_r)
+
     return nothing
 end
 

test/Project.toml

@@ -1,3 +1,5 @@
 [deps]
+CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
+DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

test/data/case5/.gitignore

@@ -1,3 +1,4 @@
 # Track inflow binary files:
 !inflow.hdr
-!inflow.bin
+!inflow.bin
+!inflow.csv