Merge pull request #509 from probcomp/20230620pfstate-copy

add Base.copy(::ParticleFilterState)

Author: ztangent

src/inference/particle_filter.jl

@@ -23,6 +23,28 @@ mutable struct ParticleFilterState{U}
     parents::Vector{Int}
 end
 
+function Base.copy(state::ParticleFilterState{U}) where U
+    return ParticleFilterState{U}(
+        Base.copy(state.traces),
+        Base.copy(state.new_traces),
+        Base.copy(state.log_weights),
+        state.log_ml_est,
+        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,
+                             hash(state.log_ml_est, 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