def compute_JSD(mu_P, Sigma_P, mu_Q, Sigma_Q, N_MC_samples):
# Jensen-Shannon distance (JSD)
out_JSD = np.empty_like(mu_P)
for index in range(mu_P.size):
RV_p = MVN(mean=mu_P[index], cov=Sigma_P[index, index])
RV_q = MVN(mean=mu_Q[index], cov=Sigma_Q[index, index])
x = RV_p.rvs(N_MC_samples)
p_x = RV_p.pdf(x)
q_x = RV_q.pdf(x)
m_x = (p_x + q_x) / 2.
y = RV_q.rvs(N_MC_samples)
p_y = RV_p.pdf(y)
q_y = RV_q.pdf(y)
m_y = (p_y + q_y) / 2.
dKL_pm = np.log(p_x / m_x).mean()
dKL_qm = np.log(q_y / m_y).mean()
out_JSD[index] = 0.5 * (dKL_pm + dKL_qm)
return out_JSD
class Q0(Q0_Base):
""" Define prior """
def __init__(self):
""" Define prior pdfs over stiffness, damping, and
noise std.
"""
self.p_k = Normal_PDF(mean=3, cov=0.5)
self.p_c = Gamma_PDF(a=1, scale=0.1)
self.p_sigma = Gamma_PDF(a=1, scale=0.1)
def logpdf(self, x):
# Convert to 2D array if currently 1D
if len(np.shape(x)) == 1:
x = np.array([x])
# Calculate logpdf
logpdf = (self.p_k.logpdf(x[:, 0]) + self.p_c.logpdf(x[:, 1]) +
self.p_sigma.logpdf(x[:, 2]))
return logpdf
def rvs(self, size):
k = np.vstack(self.p_k.rvs(size))
c = np.vstack(self.p_c.rvs(size))
sigma = np.vstack(self.p_sigma.rvs(size))
return np.hstack([k, c, sigma])
class Q0(Q0_Base):
""" Define initial proposal """
def __init__(self):
self.pdf = Normal_PDF(mean=np.zeros(2), cov=np.eye(2))
def logpdf(self, x):
return self.pdf.logpdf(x)
def rvs(self, size):
return self.pdf.rvs(size)
def sample_motion(n_samples, indices, motion_df, img, xweight=None, yweight=None):
"""
Sampel x and y motion for multivariate normal distribution.
Parameters:
- - - - -
n_samples: int
number of samples to generate
indices: int, array
indices to apply motion to
motion_df: Pandas data frame
data frame of sampled motion to fit 6d-Gaussian to
expected column names: 'xrad', 'yrad', 'zrad'
'xshift', 'yshift', 'zshift'
img: float, array
input image
xweight: float
weight to apply to xmotion
yweight: float
weight to apply to ymotion
"""
# get input image dimensions
[n, p] = img.shape
# get weights for each axis of motion
if not xweight:
xweight = 1
if not yweight:
yweight = 1
# fit distribution from which to sample motion from
mu = motion_df.mean(0)
cov = motion_df.cov()
gauss = MVN(mean=mu, cov= cov)
# sample positions in K-space tracjectory for when
# motion occurs in X and Y directions
inds = np.random.choice(indices, size=n_samples, replace=False)
inds = np.column_stack(np.unravel_index(inds, (n, p)))
# sample motion for each coordinate
samples = pd.DataFrame(
gauss.rvs(size=n_samples),
columns=motion_df.columns)
x_motion = np.zeros((n, p))
x_motion[inds[:, 0], inds[:, 1]] = xweight*np.asarray(samples['xshift'])
x_coords = np.ravel_multi_index(np.where(x_motion != 0), x_motion.shape)
y_motion = np.zeros((n, p))
y_motion[inds[:, 0], inds[:, 1]] = yweight*np.asarray(samples['yshift'])
y_coords = np.ravel_multi_index(np.where(y_motion != 0), y_motion.shape)
return [x_motion, x_coords, y_motion, y_coords]