Sampling BIP with multiple-input forward models

In many real-world inverse problems, a forward model depends on several fundamentally distinct physical quantities or parameters, rather than a single, homogeneous set of variables. For instance, in a mechanical problem, one might simultaneously infer material stiffness, boundary conditions, and external loads. These parameters often have different physical units, vastly distinct numerical scales, and completely different prior beliefs. While mathematically possible, grouping these distinct quantities into a single long vector can obscure their individual semantic properties, lead to poor inference, and make constructing customized, efficient Markov Chain Monte Carlo (MCMC) algorithms exceedingly difficult.

CUQIpy provides an elegant and mathematically rigorous solution to this challenge through its native support for multiple-input models and Gibbs sampling strategy. By allowing users to define separate geometries and prior distributions for each input element/group, developers can build a multiple-input forward Model that respects the individual nature of each input element/group. When paired with the HybridGibbs sampler, this allows users to apply tailored, element/group-wise sampling strategies, such as using a Metropolis-Hastings (MH) step for lower-dimensional groups where gradient information is difficult to obtain, while employing the Metropolis-Adjusted Langevin Algorithm (MALA) for high-dimensional ones where gradient is available.

To intuitively demonstrate this capability, we will explore a very simplified “bathtub” mixing problem. Imagine a bathtub filled with mixed water where we only have measurements of the final combined volume and the final stabilized temperature. Our goal is to solve the corresponding inverse problem: inferring the initial volume and temperature of both the hot water and the cold water sources that were used to fill the tub.

1. The forward model¶

Let $h_v$ be the volume of hot water, $h_t$ be the temperature of hot water, $c_v$ be the volume of cold water, and $c_t$ be the temperature of cold water. The forward map determines the final volume $V$ and the final temperature $T$ of the mixed water:

1. The final volume is simply the sum of the hot and cold water volumes:

   $$V = h_v + c_v$$

2. The final temperature is calculated as the weighted average of the two temperatures, based on their respective volumes:

   $$T = \frac{h_v h_t + c_v c_t}{h_v + c_v}$$

Advanced MCMC algorithms, such as the MALA or HMC, rely on gradient information to efficiently explore the target distribution. Specifically, they need the gradient of the log-posterior, which in turn requires the derivative of the forward model with respect to each input parameter.

By supplying explicit analytical gradients for $h_v$, $h_t$, $c_v$, and $c_t$, we make our sampler both significantly faster and numerically more stable compared to relying on finite-difference approximations.

Below, we define the action of the Jacobians (the directional derivatives) for each input parameter:

In CUQIpy, geometries govern how parameters are interpreted and presented. Because we have four distinct inputs with completely different meanings (two are volumes, two are temperatures), it is natural to represent each of them as a Discrete geometry. Then we define a tuple for our domain_geometry consisting of the 4 Discrete geometries. Similarly, the range_geometry defines the physical structure of our output space: exactly two discrete point measurements (temperature and volume).

Now we wrap everything together into a cuqi.model.Model object. The constructor simply requires:

1. forward: our Python function simulating the mixture forward_map.

2. gradient: the tuple of our four analytically computed gradients.

3. domain_geometry: our tuple defining the four input parameters.

4. range_geometry: our geometry specifying the two outputs.

When initialized, this object internally coordinates evaluating subsets of inputs and parsing gradients on demand during the MCMC sampling process.

2. Prior distributions¶

In real-world applications, prior distributions should be carefully elicited based on physical constraints and expert knowledge. However, for the purpose of this pedagogical example, we choose simple, somewhat arbitrary priors to illustrate the workflow in CUQIpy. Specifically, for the fluid volumes and the hot water temperature, we utilize Uniform distributions to reflect simple bounds. For the cold water temperature, we apply a TruncatedNormal distribution to demonstrate how more informative priors with specific boundaries can be incorporated.

3. Data, joint and posterior distribution¶

Then we construct the following objects:

• Data Distribution: We assume the measurements of volume and temperature are corrupted by independent Gaussian noise. The Model we defined serves as the mean of this distribution.

• Joint Distribution: We combine the data distribution and the priors into a single joint distribution.

• Posterior: By conditioning the joint distribution on our observed data, we obtain the posterior distribution of our unknown parameters.

4. Sampling with HybridGibbs¶

We will use CUQIpy’s HybridGibbs sampler. Gibbs sampling involves iteratively drawing samples from the conditional distribution of each parameter given the current state of the others. To do this in CUQIpy, we must define a sampling strategy: a dictionary that maps each of our target distributions to a specific MCMC sampler.

Here, we assign the MH sampler for the first 3 fields and the gradient-aware MALA sampler for the last field. As with the priors, this configuration is chosen primarily to demonstrate the flexibility of CUQIpy’s HybridGibbs framework which allows mixed sampling strategies. In a demanding real-world scenario, the choice of samplers would be carefully matched to the dimensionality and gradient availability for each parameter.

Then create the sampler and sample the posterior distribution.

Finally, we can visualize our samples. Because we are trying to infer 4 initial state parameters from only 2 macroscopic measurements, the inverse problem is highly underdetermined. As we can see, the temperature and volume of the mix computed using the posterior means match the true solution very well.