Adapting this Bayesian localization code to a 3-dimensional solution

I’m looking at this 2D solution to the Time-Difference-of-Arrival (TDoA) problem, using Bayesian inference and pymc3. I’ll give the abridged version here:

We have a network of n receivers, and some signal source. We know the location of each receiver in the network, and each records the Time-of-Arrival (ToA) of the signal at that receiver. We want to use this information to determine the coordinates of the signal source.

This is a classic multilateration problem and, in 2D, the simplest way to solve it is to draw the locus of possible source locations for each receiver and search for where they overlap. However, when the ToA measurements are noisy, and the speed of the signal is uncertain, the loci will not always overlap and this method will be ambiguous.

Instead, we represent the problem using Baye’s theorem: Where the location of the source x, signal speed v, and observed arrival times t are all random variables. p(x, v|t) is the posterior distribution of the source location and signal speed given the arrival times, and this is what we want to obtain. The data likelihood p(t|x,v) is just the forward physics model plus observational noise and p(x), p(v) define prior distributions over the source location and signal speed, which we must choose.

The difficulty in obtaining the posterior is in computing the evidence p(t), which is intractable for many Bayesian inference problems. Instead, we obtain samples from the posterior using the NUTS sampler in pymc3.

Step 1: Define the forward (physics) model

import numpy as np
import pymc3 as pm
import matplotlib
import matplotlib.pyplot as plt

#stations = [array of receiver coordinates] (to be defined below)
#t_obs = [array of actual observed TDOA measurements from the receivers] (to be defined below)

with pm.Model():

# Priors
x = pm.Uniform("x", lower=0, upper=1000, shape=2)# prior on the source location (m)
v = pm.Normal("v", mu=346, sigma=20)# prior on the wave speed (m/s)
t1 = pm.Uniform("t1", lower=-2, upper=4)# prior on the time offset (s)

# Physics
d = pm.math.sqrt(pm.math.sum((stations - x)**2, axis=1))# distance between source and receivers
t0 = d/v# time of arrival (TOA) of each receiver
t = t0-t1# time difference of arrival (TDOA) from the time offset

# Observations
Y_obs = pm.Normal('Y_obs', mu=t, sd=0.05, observed=t_obs)# we assume Gaussian noise on the TDOA measurements

Step 2: Simulate some observations to test our model with

# generate some test observations
N_STATIONS = 4
np.random.seed(1)
stations = np.random.randint(250,750, size=(N_STATIONS,2))# station positions (m)
x_true = np.array([500,500])# true source position (m)
v_true = 346.# speed of sound (m/s)
t1_true = 1.0# can be any constant, used to calculate TDOA values (s)
d_true = np.linalg.norm(stations-x_true, axis=1)
t0_true = d_true/v_true# true time of flight values
t_obs = t0_true-t1_true# true time difference of arrival values
np.random.seed(1)
t_obs = t_obs+0.05*np.random.randn(*t_obs.shape)# noisy observations
print("t_obs (s):", t_obs)

# t_obs and stations are passed to the pymc3 model above

Step 3:

And then we sample from the posterior of our model using:

# Sample posterior
trace = pm.sample(draws=2000,
tune=2000,
chains=4,
target_accept=0.95,

pymc3 should auto-assign the NUTS sampler, which draws samples from the posterior and stores them in the trace object. All we need to do is plot these samples to obtain an estimate of the location with its associated uncertainty, and we are done:

# get means and standard deviations of posterior samples
summary = pm.summary(trace)
mu = np.array(summary["mean"])
sd = np.array(summary["sd"])

# plot
plt.figure(figsize=(5,5))
plt.scatter(x_true, x_true, s=40, label="True source position")
ell = matplotlib.patches.Ellipse(xy=(mu, mu),
width=4*sd, height=4*sd,
angle=0., color='black', label="Posterior ()", lw=1.5)
ell.set_facecolor('none')