Replace reservoir array with reservoir trajectory

src/components/trajectories/reservoir_trajectory.jl

@@ -0,0 +1,34 @@
+export ReservoirTrajectory
+
+using MacroTools: @forward
+using Random
+
+mutable struct ReservoirTrajectory{B,R<:AbstractRNG}  <: AbstractTrajectory
+    buffer::B
+    n::Int
+    capacity::Int
+    rng::R
+end
+
+@forward ReservoirTrajectory.buffer Base.keys, Base.haskey, Base.getindex
+
+Base.length(x::ReservoirTrajectory) = length(x.buffer[1])
+
+function ReservoirTrajectory(capacity, kw::Pair{Symbol, DataType}...;n=0, rng=Random.GLOBAL_RNG)
+    buffer = Trajectory(;(s => Vector{t}() for (s,t) in kw)...)
+    ReservoirTrajectory(buffer, n, capacity, rng)
+end
+
+function Base.push!(b::ReservoirTrajectory; kw...)
+    b.n += 1
+    if b.n <= b.capacity
+        push!(b.buffer; kw...)
+    else
+        i = rand(b.rng, 1:b.n)
+        if i <= b.capacity
+            for (k,v) in kw
+                b.buffer[k][i] = v
+            end
+        end
+    end
+end

src/components/trajectories/trajectories.jl

@@ -1,2 +1,3 @@
 include("abstract_trajectory.jl")
 include("trajectory.jl")
+include("reservoir_trajectory.jl")

src/utils/device.jl

@@ -15,20 +15,20 @@ send_to_device(::Val{:cpu}, x::CuArray) = adapt(Array, x)
 send_to_device(::Val{:gpu}, x) = Flux.fmap(a -> adapt(CuArray{Float32}, a), x)
 
 const KnownArrayVariants = Union{
-    SubArray{<:Any,<:Any,<:Union{ReservoirArrayBuffer,CircularArrayBuffer,ElasticArray}},
+    SubArray{<:Any,<:Any,<:Union{CircularArrayBuffer,ElasticArray}},
     Base.ReshapedArray{
         <:Any,
         <:Any,
         <:SubArray{
             <:Any,
             <:Any,
-            <:Union{ReservoirArrayBuffer,CircularArrayBuffer,ElasticArray},
+            <:Union{CircularArrayBuffer,ElasticArray},
         },
     },
     Base.ReshapedArray{
         <:Any,
         <:Any,
-        <:Union{ReservoirArrayBuffer,CircularArrayBuffer,ElasticArray},
+        <:Union{CircularArrayBuffer,ElasticArray},
     },
     SubArray{
         <:Any,
@@ -39,7 +39,7 @@ const KnownArrayVariants = Union{
             <:SubArray{
                 <:Any,
                 <:Any,
-                <:Union{ReservoirArrayBuffer,CircularArrayBuffer,ElasticArray},
+                <:Union{CircularArrayBuffer,ElasticArray},
             },
         },
     },

src/utils/utils.jl

@@ -1,6 +1,5 @@
 include("printing.jl")
 include("base.jl")
 include("circular_array_buffer.jl")
-include("reservoir_array_buffer.jl")
 include("device.jl")
 include("sum_tree.jl")

test/components/trajectories.jl

@@ -207,4 +207,41 @@
         empty!(t)
         @test length(t[:state]) == 0
     end
+
+    @testset "ReservoirTrajectory" begin
+        # test length
+        t = ReservoirTrajectory(3, :a=>Array{Float64,2}, :b=>Bool)
+        push!(t;a=rand(2,3),b=rand(Bool))
+        @test length(t) == 1
+        push!(t;a=rand(2,3),b=rand(Bool))
+        @test length(t) == 2
+        push!(t;a=rand(2,3),b=rand(Bool))
+        @test length(t) == 3
+
+        for _ in 1:100
+            push!(t;a=rand(2,3),b=rand(Bool))
+        end
+
+        @test length(t) == 3
+
+        # test distribution
+
+        Random.seed!(110)
+        k, n, N = 3, 10, 10000
+        stats = Dict(i => 0 for i in 1:n)
+        for _ in 1:N
+            t = ReservoirTrajectory(k, :a=>Array{Int, 2}, :b=>Int)
+            for i in 1:n
+                push!(t;a=i .* ones(Int, 2, 3), b=i)
+            end
+
+            for i in 1:length(t)
+                stats[t[:b][i]] += 1
+            end
+        end
+
+        for v in values(stats)
+            @test isapprox(v/N, k/n;atol=0.03)
+        end
+    end
 end

test/utils/utils.jl

@@ -1,4 +1,3 @@
 include("base.jl")
 include("circular_array_buffer.jl")
-include("reservoir_array_buffer.jl")
 include("device.jl")