hash, eq for PFState; tests

Author: georgematheos

src/inference/particle_filter.jl

@@ -24,14 +24,28 @@ mutable struct ParticleFilterState{U}
 end
 
 function Base.copy(state::ParticleFilterState{U}) where U
-    ParticleFilterState{U}(
-        copy(state.traces),
-        copy(state.new_traces),
-        copy(state.log_weights),
+    return ParticleFilterState{U}(
+        Base.copy(state.traces),
+        Base.copy(state.new_traces),
+        Base.copy(state.log_weights),
         state.log_ml_est,
-        copy(state.parents)
+        Base.copy(state.parents)
     )
 end
+function Base.:(==)(a::ParticleFilterState{U}, b::ParticleFilterState{V}) where {U, V}
+    return U == V &&
+        a.traces == b.traces &&
+        a.log_weights == b.log_weights &&
+        a.log_ml_est == b.log_ml_est &&
+        a.parents == b.parents
+end
+function Base.hash(state::ParticleFilterState{U}, h::UInt) where {U}
+    return hash(U,
+    hash(state.traces,
+    hash(state.log_weights .+ (state.log_ml_est - log(length(state.log_weights))),
+    hash(state.parents, 
+         h))))
+end
 
 """
     traces = get_traces(state::ParticleFilterState)

test/inference/particle_filter.jl

@@ -169,4 +169,31 @@ end
     @test isapprox(expected_log_ml, actual_log_ml_est, atol=0.02)
 end
 
+@testset "PF state" begin
+    @gen function _foo()
+        x ~ normal(0, 1)
+    end
+    st = Gen.ParticleFilterState{Gen.Trace}(
+        [simulate(_foo, ()) for _=1:10],
+        Vector{Gen.Trace}(undef, 10),
+        [0. for _=1:10],
+        0.,
+        collect(1:10)
+    )
+    st2 = copy(st)
+    @test st == st2
+    @test hash(st) == hash(st2)
+    st2.log_weights[1] = 1.
+    @test st.log_weights[1] == 0.
+    @test st != st2
+    @test hash(st) != hash(st2)
+
+    # test that the other fields are independent copies too:
+    st.traces[1] = simulate(_foo, ())
+    @test st.traces != st2.traces
+    st.log_ml_est = 1.
+    @test st.log_ml_est == 1.
+    st.parents[1] = 5
+    @test st2.parents[1] == 1
 end
+end