def sample(self, Z):
"When real=False we assume Z[0] = Z.real and Z[1] = Z.imag"
if not self.real:
Z = array2complex(Z)
X = fftn(Z, norm="ortho")
X = complex2array(X)
return X
def sample(self, Z):
"We assume Z[0] = Z.real and Z[1] = Z.imag"
Z = array2complex(Z)
X = self.W.dot(Z)
X = complex2array(X)
assert X.shape == (2, self.Nx)
return X
def compute_forward_mean(self, az, bz, ax, bx):
# estimate x from x = Wz we have rx = W rz
rz = self.compute_backward_mean(az, bz, ax, bx)
rz = array2complex(rz)
rx = self.W @ rz
rx = complex2array(rx)
return rx
def compute_backward_message(self, az, bz, ax, bx):
# z = U.conj().T x
az_new = ax
bx = array2complex(bx)
bz_new = self.U.conj().T @ bx
bz_new = complex2array(bz_new)
return az_new, bz_new
def compute_forward_message(self, az, bz, ax, bx):
# x = U z
ax_new = az
bz = array2complex(bz)
bx_new = self.U @ bz
bx_new = complex2array(bx_new)
return ax_new, bx_new
def compute_forward_message(self, az, bz, ax, bx):
# x = FFT z
ax_new = az
if not self.real:
bz = array2complex(bz)
bx_new = fftn(bz, norm="ortho")
bx_new = complex2array(bx_new)
return ax_new, bx_new
def compute_log_partition(self, az, bz, ax, bx):
rz = self.compute_backward_mean(az, bz, ax, bx)
b = complex2array(
array2complex(bz) + self.W.conj().T @ array2complex(bx)
)
a = az + ax * self.spectrum
logZ = 0.5 * np.sum(b * rz) + np.sum(np.log(2 * np.pi / a))
return logZ
def sample(self, Z):
"When real=False we assume Z[0] = Z.real and Z[1] = Z.imag"
if not self.real:
Z = array2complex(Z)
X = self.convolve(Z)
if not self.real:
X = complex2array(X)
return X
def compute_backward_message(self, az, bz, ax, bx):
# z = IFFT x
az_new = ax
bx = array2complex(bx)
bz_new = ifftn(bx, norm="ortho")
if self.real:
bz_new = np.real(bz_new)
else:
bz_new = complex2array(bz_new)
return az_new, bz_new
def compute_forward_mean(self, az, bz, ax, bx):
# estimate x from x = Wz we have rx = W rz
rz_fft = self.compute_backward_mean(az, bz, ax, bx, return_fft=True)
rx_fft = self.w_fft * rz_fft
rx = ifftn(rx_fft)
if self.real:
rx = np.real(rx)
else:
rx = complex2array(rx)
return rx
def compute_backward_mean(self, az, bz, ax, bx):
# estimate z from x = Wz
bz = array2complex(bz)
bx = array2complex(bx)
if self.precompute_svd:
bx_svd = self.U.conj().T @ bx
bz_svd = self.V.conj().T @ bz
resolvent = 1 / (az + ax * self.spectrum)
rz_svd = resolvent * (bz_svd + self.S.conj().T @ bx_svd)
rz = self.V @ rz_svd
else:
a = az * np.identity(self.Nz) + ax * self.C
b = (bz + self.W.conj().T @ bx)
rz = np.linalg.solve(a, b)
rz = complex2array(rz)
return rz
def compute_backward_mean(self, az, bz, ax, bx, return_fft=False):
# estimate z from x = Wz
if not self.real:
bz = array2complex(bz)
bx = array2complex(bx)
bx_fft = fftn(bx)
bz_fft = fftn(bz)
resolvent = 1 / (az + ax * self.spectrum)
rz_fft = resolvent * (bz_fft + self.w_fft_bar * bx_fft)
if return_fft:
return rz_fft
rz = ifftn(rz_fft)
if self.real:
rz = np.real(rz)
else:
rz = complex2array(rz)
return rz
def _test_function_posterior(self, likelihood, records, places=12):
for record in records:
az, bz = record["az"], record["bz"]
y = float(likelihood.y)
if isinstance(likelihood, ModulusLikelihood):
rz, vz = explicit_complex_integral(az, bz, y, likelihood)
bz = complex2array(np.array(bz))
rz_hat, vz_hat = likelihood.compute_backward_posterior(
az, bz, y)
rz_hat = array2complex(rz_hat)
rz_hat = complex(rz_hat)
else:
rz, vz = explicit_real_integral(az, bz, y, likelihood)
rz_hat, vz_hat = likelihood.compute_backward_posterior(
az, bz, y)
rz_hat = float(rz_hat)
msg = f"record={record} likelihood={likelihood}"
self.assertAlmostEqual(rz, rz_hat, places=places, msg=msg)
self.assertAlmostEqual(vz, vz_hat, places=places, msg=msg)
def compute_log_partition(self, az, bz, ax, bx):
b = complex2array(
array2complex(bz) + self.U.conj().T @ array2complex(bx))
a = az + ax
logZ = 0.5 * np.sum(b**2 / a) + self.N * np.log(2 * np.pi / a)
return logZ
