How to condition a causal model on a time-series of observations

I’m learning how to use Pyro for variational inference (while learning that too).

I’ve graduated from a simple example with N independent trials and now I want to do inference on a time-series (trajectory of a discrete-time stochastic system).

I have a hybrid model with both continuous and discrete latent random variables:

def random_shock(scale_guess, b_guess, epsilon_guess):
    # Scale and b are always positive
    scale = pyro.sample("scale", dist.LogNormal(np.log(scale_guess), 0.25))
    b = pyro.sample("b", dist.LogNormal(np.log(b_guess), 0.25))
    epsilon = pyro.sample("epsilon", dist.LogNormal(np.log(epsilon_guess), 0.25))
    alpha = pyro.sample("alpha", dist.Bernoulli(epsilon))
    alpha = alpha.long()
    wp_dist = dist.Normal(torch.tensor([0.0, 0.0])[alpha], 
                          torch.tensor([scale, scale*b])[alpha])
    measurement = pyro.sample('obs', wp_dist)
    return measurement

def infrequent_step_5(scale_guess, b_guess, epsilon_guess):
    y = 0
    measurements = []
    for k in range(5):
        y = y + random_shock(scale_guess, b_guess, epsilon_guess)
    return torch.tensor(measurements)

epsilon_guess = 0.05
b_guess = 50
scale_guess = 0.03

# Draw random samples
for _ in range(5):
    print(infrequent_step_5(scale_guess, b_guess, epsilon_guess))

Output seems to be as expected:

tensor([0.0156, 0.0179, 0.0519, 0.0739, 0.0711])
tensor([ 1.7345e-04, -4.2217e-03, -6.6964e-05,  1.6579e-02, -3.5251e-02])
tensor([ 0.0012, -0.0081,  0.0535, -0.2341, -0.1900])
tensor([ 0.0189, -0.0123,  0.0166,  0.0029, -0.0133])
tensor([-0.0015,  0.0066,  0.0482,  0.1065,  0.1311])

Now, I’m stuck on how to make a model conditioned on a set of observations and how to set up the SVI optimization.

Here is my attempt at the conditioned model which is based on this example from the documentation:

def infrequent_step_5_conditioned(scale_guess, b_guess, epsilon_guess, observations):
    y = 0  # initial value
    for k, y_m in enumerate(observations):
        y = y + random_shock(scale_guess, b_guess, epsilon_guess)
        pyro.sample(f"y_{k}", y, obs=y_m)

My questions are:

  • is infrequent_step_5_conditioned correct?
  • how do I do SVI with this model and a set of observations? (links to relevant examples appreciated)
  • is it possible to condition the model on multiple observed trajectories? I guess this would involve a with pyro.plate ... construct.