Reverse Jump

Caution

This feature requires more testing.

We now solve a classical linear regression problem.

Note

We explain some snippets and the path of complete code is examples/rj_sampling.py.

The signal consists of \(100\) samples and is composed of noise and a linear function to be estimated.

\[ s = n + h, \]

where \(n\) is Gaussian white noise with \(\sigma=0.1\) and \(h(x)=1+0.05x,\;x\in[0,\,1]\).

Assuming that we do not know the optimal degree of polynomial to fit the data. We will choose from degree \(0\) and \(1\) polynomials:

\[\begin{split} h = \begin{cases} a_0, & \mathrm{degree0}; \\ a_0 + a_1 x, & \mathrm{degree1}. \end{cases} \end{split}\]

# ============== Generate dataset ================
nsample = 100
noise_sigma = .1
x = np.linspace(0, 1, nsample)
nbyte = nsample * np.dtype(np.float64).itemsize if iproc == 0  else 0
win = MPI.Win.Allocate_shared(nbyte, comm=comm)
shared_data, _ = win.Shared_query(0)
signal = np.ndarray(shape=nsample, dtype=np.float64, buffer=shared_data)
win.Fence()
if iproc == 0:
    noise = ts.random.rng.normal(0., noise_sigma, nsample)
    hsignal = 1. + .05 * x
    signal[:] = noise + hsignal

    import matplotlib.pyplot as plt
    plt.plot(x, hsignal, color="tab:red", label="$h$")
    plt.plot(x, signal, label="$s$")
    plt.xlabel("x")
    plt.legend()
    plt.savefig(TMP_DIR/"rj_sampling_signal.pdf", dpi=150, bbox_inches="tight")

We generate the simulated \(s\) in single process since all the processes should “see” identical signal. Other processes can obtain it by bcast operation in principle. Here, we show another way: shared window. Unlike fork in Unix system, processes in mpi program do not share identical variables even if there is no modication. In this example, we use only 8 processes and have known in advance that all of these are located in one node. So, we only store one version of \(s\) in physical memory of this node and all the process can access it with the shared window. This method will be particularly useful for giant dataset.

win.Fence()

# ============== Prior ================
priors = {
    "degree0": (.5, {"a0": ts.prior.Uniform(.5, 1.5)}),
    "degree1": (.5, {"a0": ts.prior.Uniform(.5, 1.5),
                 "a1": ts.prior.Uniform(-.2, .2)})
}
priors = ts.prior.ModelPrior(priors)

# ============== Likelihood ================
class ToyLikelihood:
    def calc_log_likelihood(self, state):
        h_hypo = state.params[0]
        if state.model == "degree1":
            h_hypo += state.params[1] * x
        residual = signal - h_hypo
        ll = -.5 * np.sum(residual**2) / noise_sigma**2 - nsample * np.log(TWOPI**.5 * noise_sigma)    
        return ll

priors consists of prior information of each model. We need to set the probability (density) of the model and parameters of the model. Here, the probabilities of “degree0” and “degree1” are both \(0.5\).

toylikelihood = ToyLikelihood()

# ============== Move ================
within_move = {
    "degree0": ts.move.GaussianMHMove(np.array([.05])**2),
    "degree1": ts.move.GaussianMHMove(np.array([.05, .01])**2)
}
rjmove0 = ts.move.RJMove(["degree1"], 
                         add_params=[ts.prior.Prior(a1=ts.prior.Uniform(-.2, .2))], 
                        )
rjmove1 = ts.move.RJMove(["degree0"])
between_move = {"degree0": rjmove0, "degree1": rjmove1}
move = ts.move.ModelMove(within_move, between_move)

We set the method of move of individual model and gather them into within_move. rjmove is treated with a similar way. Notice that the first parameter is the model to which can be jumped from current model.

The result of the parameter estimation is

_ = result.plot_dist_model()

_ = result.plot_dist_params("degree1", truths=[1., .05])