Example - Spring Mass System — SMCPy 1.0 documentation

# Example - Spring Mass System

This example provides a simple demonstration of SMCPy functionality. The goal is to inversely determine the uncertainty in the model parameters given a set of noisy observations of the spring mass system using Sequential Monte Carlo (SMC) and to compare the results to Markov chain Monte Carlo (MCMC). As SMC relies on a MCMC kernel, a fully functional MCMC sampler is included in the SMCPy package. In this example, it is used for verification of the SMC sampler. The example covers all steps for implementing SMC and MCMC samplers using SMCPy, including the creation of a user-defined computational model (spring mass numerical integrator) that uses the standardized SMCPy interface, defining prior distributions, initializing the samplers, computing statistical moments with the resulting estimators, and plotting the results. The full source code for this example can be found in the SMCPy repository: `/SMCPy/examples/spring_mass/spring_mass_example.py` and `/SMCPy/examples/spring_mass/mcmc_verify/spring_mass_mcmc.py` for the SMC and MCMC samplers, respectively. Data generation was conducted using `/SMCPy/examples/spring_mass/generate_noisy_data.py`.

## Problem Specification

The governing equation of motion for the system is given by

(1)\[m_s \ddot{z} = -k_s z + m_s g\]

where \(m_s\) is the mass, \(k_s\) is the spring stiffness, \(g\) is the acceleration due to gravity, \(z\) is the vertical displacement of the mass, and \(\ddot{z}\) is the acceleration of the mass. The true stiffness and gravitational constant are unknown. The goal of this example is to, by observing the motion of the mass over time, \(z_t\), estimate both \(K_s\) and \(G\), which are random variables representing the uncertainty in the value of spring stiffness and the gravitational constant, respectively. For reasons outside the scope of this example, the problem becomes ill-posed unless the stiffness is normalized by mass such that \(k_s^* = k_s/m_s\) and \(k_s^*\) are realizations of the random variable of interest \(K_s^*\). The equation of motion given by Equation (1) is now

(2)\[\ddot{z} = -k_s^* z + g.\]

Given observations \(z_t\) which, assuming measurement noise exists, are realizations of the random variable \(Z_t\), the relationship between the computational model, measurement noise, and observations is

(3)\[Z_t = f_t(K_s^*, G) + \epsilon_t\]

Here, \(\epsilon_t\) is the measurement noise associated with each observation of displacement taken over time (also a random variable), and \(f_t\) is the computational model response at time \(t\), which involves numerical integration of Equation (2).

The inverse solution of equations in the form of (2) is the posterior distribution, or, in the context of this example, the joint distribution of \(K_s^*\) and \(G\) conditional on the observed data, \(z_t\). While difficult to solve directly, the posterior distribution can be approximated via sampling methods such as SMC and MCMC. This example covers this process using the SMCSampler and MCMCSampler classes in the SMCPy Python module. In particular, the objective is to approximate the expected value of the model parameters given observations \(z_t\).

The MCMC sampler draws samples from the unknown posterior distribution by forming a Markov chain through the parameter space whose stationary distribution is the posterior. The samples forming the chain can then be used to build an estimator of the expected value

\[E[X] = \frac{1}{N} \sum_{i=1}^N X_i\]

where \(X_i\) is the random variable of interest, and \(i=1,\ldots,N\), with \(N\) being the number of equally-weighted samples drawn using the MCMC sampler.

While MCMC is a proven approach to evaluting the quantities of interest, it can be slow if the computational model is expensive. The Markovian nature of MCMC means it is inherently a serial process. SMC, on the other hand, is a parallelizable alternative that uses weighted samples called “particles.” In SMC, a sequence of target distributions is defined that gradually transitions from the typically-known prior distribution to the posterior distribution of interest. A particle filtering framework based on importance sampling and a MCMC transition kernel can then be introduced that allows for recursive estimation of each target in the sequence. Taking the final particle state, the SMC estimator is

\[E[X] = \sum_{j=1}^M W_j X_j\]

where \(W_j\) is the normalized weight associated with the \(j^{th}\) particle and \(M\) is the total number of particles.

For this example, synthetic data will be generated by adding noise to a model response for a given set of “true” parameters. Next, MCMC and SMC will be used to estimate the expected value of these parameters with the expectation that these estimates are close to the true parameter values. Note that the code used to generate results in this example is actually split between two files in the SMCPy package. The two will be combined here and redundant code will be skipped.

## Step 1: Initialization; define the model and generate the data

Begin by importing the needed Python modules, including SMCPy sampler classes and the SpringMassModel class that defines the spring mass numerical integrator:

    import numpy as np
    from smcpy.smc.smc_sampler import SMCSampler
    from smcpy.mcmc.mcmc_sampler import MCMCSampler
    from spring_mass_models import SpringMassModel

Below is a snippet of the SpringMassModel class; the entire class can be found in the SMCPy repo (`/SMCPy/examples/spring_mass/spring_mass_model.py`):

    from smcpy.model.base_model import BaseModel

    ...

    class SpringMassModel(BaseModel):
        '''
        Defines Spring Mass model with 2 free params (spring stiffness, k &
        mass, m)
        '''
        def __init__(self, state0=None, time_grid=None):

Note that user-defined models in SMCPy must inherit from the SMCPy abstract class `BaseModel` and implement an `evaluate` function that accepts and returns numpy arrays for inputs and outputs, respectively. Here, the `state0` argument defines the initial state of the spring mass system, and `time_grid` defines the times at which to return displacement.

The first step in an analysis is to obtain data from which to make an inference. In this example, this data will come in the form of observations of the z-displacement of the mass made over time. For demonstration purposes, the data will be generated from the spring mass model, and noise will be added by sampling from a zero-mean Gaussian distribution and adding these values to the model output. While not a realistic case, it is typical to generate and use synthetic data in this manner for verification purposes when performing inverse uncertainty quantification. The following code snippet is from `generate_noisy_data.py` and demonstrates the generation of data assuming the ground truth parameters, or those that are unknown and are to be estimated, are \(k_s^*=1.67\) and \(g^*=4.62\):

    # Initialize model
    state0 = [0., 0.]                        #initial conditions
    measure_t_grid = np.arange(0., 5., 0.2)  #time
    model = SpringMassModel(state0, measure_t_grid)

    # Define the ground truth
    true_params = {'K': 1.67, 'g': 4.62}

    # Load data
    noise_stddev = 0.5
    displacement_data = model.generate_noisy_data_with_model(noise_std_dev, truth_params)

Once the data has been generated, the only remaining task is to use the SMC and MCMC sampler classes to generate estimates of the parameter expected values.

## Step 2: Perform Parameter Estimation using SMCPy

To instance the SMC sampler class, the data, model and parameter prior distributions must be passed to the constructor. The first two have been defined already. The parameter priors can be defined using the SMCPy standard format. In this case, both prior distributions will be defined as Uniform bounded from 0.0 to 10.0.

    # Define prior distributions
    param_priors = {'K': ['Uniform', 0.0, 10.0],
                    'g': ['Uniform', 0.0, 10.0]}

Other options, such as Normal or Truncated Normal distributions are supported.

All that is left now is to use the `sample()` method, which requires the number of particles, number of time steps (i.e., the number of target distributions to define in the sequence) and the number of MCMC proposals to make when implementing the MCMC transition kernel between targets. Another optional but useful argument used here is the effective sample size (ESS) threshold. The ESS is a measure of how many particles in the overall population have significant weight. If this number falls below the ESS threshold, the particles are resampled with replacement to reduce degeneracy.

    # SMC sampling
    num_particles = 5000
    num_time_steps = 20
    num_mcmc_steps = 1
    smc = SMCSampler(displacement_data, model, param_priors)
    pchain = smc.sample(num_particles, num_time_steps, num_mcmc_steps, noise_stddev,
                        ess_threshold=num_particles*0.5)

The object returned by the sampler is referred to as a “particle chain,” which is just a data storage object that contains all particle values, weights, and likelihoods computed at each step in the sampling process. This object has some useful methods, such as plotting tools and estimators. One such estimator is that of Equation (5), the quantity of interest for this example.

    # Calculate smc-estimated means
    smc_means = pchain.get_mean()

These values will be stored to compare with MCMC in the next section.

Before moving on, a pairwise plot can be generated to visualize parameter uncertainty and correlation using the `plot_pairwise_weights()` method of the particle chain object. The particles are represented as colored dots, with the colorbar showing normalized particle weight. The mean value of each parameter is shown as a dotted orange line.

    pchain.plot_pairwise_weights(save=True, show=False)

## Step 3: Perform Parameter Estimation using MCMCPy

Repeating the estimation of parameter means is similar with the MCMCPy submodule. Instancing the `MCMCSampler` class is done in the same way as before.

    mcmc = MCMCSampler(displacement_data, model, param_priors)

Additional parameters have to be set prior to sampling, however. In this example, it is assumed that the measurement noise variance is known, so the variance will be fixed and provided in the form of the standard deviation used when generating the data, previously. Additionally, an initial guess (i.e., where to initialize the Markov chain within the parameter space) are specified.

    initial_guess = {'K': 1.0, 'g': 1.0}
    mcmc.generate_pymc_model(q0=initial_guess, std_dev0=noise_stddev, fix_var=True)

When using the `sample()` method, the number of samples, \(N\), must be defined along with the number of those samples that will be discarded as burn-in. Burn-in is defined as the samples generated prior to the Markov chain reaching its stationary condition. In practice, the number of samples in the burn-in period is unknown and conservatively estimated.

    num_samples = 100000
    num_samples_burnin = 5000
    mcmc.sample(num_samples, num_samples_burnin)

The MCMC module is a wrapper built around the PyMC package (https://github.com/pymc-devs/pymc). The stored samples can be accessed through the PyMC trace object as follows: `<mcmc_object>.MCMC.trace(<param_name>)[:]`. The means can be calculated from the trace.

    Kmean = np.mean(mcmc.MCMC.trace('K')[:])
    gmean = np.mean(mcmc.MCMC.trace('g')[:])

A pairwise plot can also be generated with the MCMC results. Note that the samples are not weighted in this case.

    mcmc.plot_pairwise(keys=['K', 'g'], filename='mcmc_pairwise.png')

## Comparing the Results

The mean values obtained for a given run of the example scripts is shown below. Note that both SMC and MCMC sampling involve random numbers, meaning these numbers will be slightly different for every run. The true parameters set in the data generation step were \(k_s^* = 1.67\) and \(g^* = 4.62\).

Parameter | MCMC-estimated Mean | SMC-estimated mean
--- | --- | ---
\(k_s^*\) | 1.5397276686462869 | 1.5567662328730911
\(g^*\) | 4.231683157045673 | 4.276694515489001

## Running the SMCSampler in Parallel with mpi4py

As mentioned, the `SMCSampler` class was designed with high performance computing in mind. The sampler uses the mpi4py package (https://bitbucket.org/mpi4py/mpi4py) to run model evaluations at each SMC step in parallel. The SMC example script can be run in parallel using `mpirun -np <number of processors> python spring_mass_example.py`.