Enable general observation function for PDE, and accessing original solution (or other observable quantities)

[amal-ghamdi]

closes #620

[chaozg]

Hi @amal-ghamdi , thanks for implementing this feature.

Do you plan to have more demos on clarifying how this observation function could be used? For example, I wonder with an observation_map like lambda u, grid_obs: u**2, how do we compute the jacobian? I mean could the gradient be computed automatically? Is the chain product broken or 'preserved'? I guess a more common observation would be lambda u, grid:u[some indices]. Can we also compute the gradient in this case and use ULA/NUTS?

[amal-ghamdi]

Hi @amal-ghamdi , thanks for implementing this feature.

Do you plan to have more demos on clarifying how this observation function could be used? For example, I wonder with an observation_map like lambda u, grid_obs: u**2, how do we compute the jacobian? I mean could the gradient be computed automatically? Is the chain product broken or 'preserved'? I guess a more common observation would be lambda u, grid:u[some indices]. Can we also compute the gradient in this case and use ULA/NUTS?

Many thanks @chaozg for your review. I updated a howto (demos/howtos/TimeDependentLinearPDE.py) we already have to show some usage of the observation_map. The gradient is not supported automatically, the gradient of the entire forward map can be provided by the user when creating the PDEModel object. More elegant or efficient features can be done (like providing gradient for observation operator), but, I believe, it is not within the scope of this PR.

On the posterior level, one can enable FD gradient, regardless of the underlying model.

[amal-ghamdi]

Many thanks @chaozg and @jakobsj for your valuable comments! I updated the PR quite a bit to address them. I am re-requesting your review.

[chaozg]

LGTM

[jakobsj]

Changes look good and to comprehensively address the comments received. Only some very minor in line comments, so approving now.