probcomp / Gen.jl

A general-purpose probabilistic programming system with programmable inference
https://gen.dev
Apache License 2.0
1.79k stars 160 forks source link

Use map_optimize in particle filter #533

Open lwang19-ai opened 4 weeks ago

lwang19-ai commented 4 weeks ago

Is it possible to use map_optimize in the particle filter process? I did something like:

state_atm = Gen.initialize_particle_filter(unfold_atm, (start_ass_time,), init_obs_atm, num_particles)
selection = Gen.select(:chain => 1 => :u => 4 => :u)
Gen.map_optimize(state_atm.traces[1], selection)

Then I get the error: MethodError: no method matching zero(::Type{Any}) The error is from line choice_gradients in the map_optimize function. I am unsure if it's because of my implementation or if it's not supported here.

georgematheos commented 4 weeks ago

Hi @lwang19-ai, sorry you're running into this issue! Can you please share the stacktrace?

lwang19-ai commented 4 weeks ago

@georgematheos Sure.

image

# Generate proposal samples by doing pF separately.
# Then regenerate the proposal samples from first time step.

using Gen, CSV, DataFrames
include("../forward_models/gen_ks_atm_unfold.jl")

# Read from CSV files
df_atm_obs_read = CSV.read("../synthetic_truth/atm_obs.csv", DataFrame)
df_atm_ref_read = CSV.read("../synthetic_truth/atm_ref.csv", DataFrame)

# convert these back into arrays:
atm_obs = Matrix(df_atm_obs_read)

Nx = 32
Nt = 20
num_particles = 100
init_obs_atm = Gen.choicemap()
start_ass_time = 4
for i in 4:4:Nx
    init_obs_atm[(:chain=>start_ass_time=>:u=>i=>:atm_obs)] = atm_obs[start_ass_time+1, i] # Initial observation of atmosphere
end

state_atm = Gen.initialize_particle_filter(unfold_atm, (start_ass_time,), init_obs_atm, num_particles)

state_atm = Gen.map_optimize(new_trace_atm, Gen.select(:chain => 1 => :u => 1 => :u))
'''