def test_masker_simple():
mask = np.zeros((3, 2), dtype=np.bool)
mask[0, 1] = 1
mask[2, 0] = 1
x = np.arange(6).reshape((3, 2))
testing.assert_array_equal(masker(x, mask, order="F"), np.array([4, 1]))
# testing.assert_array_equal(masker(x, mask, order="C"), np.array([1, 4]))
# testing.assert_array_equal(masker(x, mask, order="C"), x[mask])
x3 = np.arange(12).reshape((3, 2, 2))
testing.assert_array_equal(masker(x3, mask, order="F"),
np.array([[8, 9], [2, 3]]))
def _forward_single_rep(self, x):
xp = self.xp
y = xp.zeros(
self.shape_in,
dtype=xp.result_type(x, xp.complex64),
order=self.order,
)
if self.use_fft_shifts:
x = self.ifftshift(x, axes=self.fftshift_axes)
if self.preplan_pyfftw:
y = self.fftn(x)
else:
y = self.fftn(x, axes=self.fft_axes)
if self.use_fft_shifts:
y = self.fftshift(y, axes=self.fftshift_axes)
if self.sample_mask is not None:
# y = y[self.sample_mask]
y = masker(
y,
mask=self.sample_mask,
order=self.order,
# mask_idx_ravel=self.sample_mask_idx,
)
return y
def test_masker():
am = masker(arr, mask, order="F", squeeze_output=True)
testing.assert_array_almost_equal(am, [0, 8, 1, 7])
assert am.shape == (nmask, )
expected_embed = arr * mask
testing.assert_array_equal(expected_embed, embed(am, mask, order="F"))
# am = masker(arr, mask, order="C", squeeze_output=True)
# testing.assert_array_almost_equal(am, [0, 1, 7, 8])
# assert am.shape == (nmask,)
# testing.assert_array_equal(expected_embed, embed(am, mask, order="C"))
am = masker(arr, mask, order="F", squeeze_output=False)
testing.assert_array_almost_equal(am[:, 0], [0, 8, 1, 7])
assert am.shape == (nmask, 1)
testing.assert_array_equal(expected_embed, embed(am, mask, order="F"))
def test_masker_multidim():
nreps = 8
arr3 = np.tile(arr[..., np.newaxis], (1, 1, nreps))
am3 = masker(arr3, mask, order="F")
assert am3.shape == (nmask, nreps)
testing.assert_array_almost_equal(am3[:, 0], [0, 8, 1, 7])
expected_embed = arr3 * mask[..., np.newaxis]
testing.assert_array_almost_equal(expected_embed,
embed(am3, mask, order="F"))
def test_TV_masked(xp):
# generate a random image-domain mask
# im_mask = xp.random.standard_normal(c.shape) > -0.2
# test 1D in/outputs but using masked values for the input array
# r = pywt.data.camera().astype(float)
r = xp.random.randn(16, 16)
x, y = xp.meshgrid(
xp.arange(-r.shape[0] // 2, r.shape[0] // 2),
xp.arange(-r.shape[1] // 2, r.shape[1] // 2),
indexing="ij",
sparse=True,
)
# make a circular mask
im_mask = xp.sqrt(x * x + y * y) < r.shape[0] // 2
# TODO: order='C' case is currently broken
for order, mask_out in product(["F"], [None, im_mask]):
r_masked = masker(r, im_mask, order=order)
Wm = TV_Operator(
r.shape,
order=order,
mask_in=im_mask,
mask_out=mask_out,
nd_input=True,
nd_output=False,
random_shift=True,
**get_loc(xp),
)
out = Wm * r_masked
if mask_out is None:
assert_(out.ndim == 1)
assert_(out.size == r.ndim * r.size)
out = out.reshape(r.shape + (r.ndim, ), order=order)
else:
assert_(out.ndim == 1)
nmask = mask_out.sum()
if xp != np:
nmask = nmask.get()
assert_(out.size == Wm.ndim * nmask)
out = embed(out, mask_out, order=order)
Wm.H * (Wm * r_masked)
def test_FiniteDifference_masked(xp, shape):
# test 1D in/outputs but using masked values for the input array
# r = pywt.data.camera().astype(float)
rstate = xp.random.RandomState(1234)
r = rstate.randn(*shape)
x, y = xp.meshgrid(
xp.arange(-r.shape[0] // 2, r.shape[0] // 2),
xp.arange(-r.shape[1] // 2, r.shape[1] // 2),
indexing="ij",
sparse=True,
)
# make a circular mask
im_mask = xp.sqrt(x ** 2 + y ** 2) < r.shape[0] // 2
# TODO: order='C' case is currently broken
for order, mask_out in product(["F"], [None, im_mask]):
r_masked = masker(r, im_mask, order=order)
Wm = FiniteDifferenceOperator(
r.shape,
order=order,
use_corners=True,
mask_in=im_mask,
mask_out=mask_out,
nd_input=True,
nd_output=mask_out is not None,
random_shift=True,
**get_loc(xp),
)
out = Wm * r_masked
if mask_out is None:
assert_(out.ndim == 1)
assert_(out.size == Wm.num_offsets * r.size)
out = out.reshape(r.shape + (Wm.num_offsets,), order=order)
else:
assert_(out.ndim == 1)
assert_(out.size == Wm.num_offsets * mask_out.sum())
out = embed(out, mask_out, order=order)
Wm.H * (Wm * r_masked)
def forward(self, x):
xp = self.xp
if x.size % self.nargin != 0:
raise ValueError("shape mismatch for forward DWT")
if (self.mask_in is not None) and (not self.nd_input):
x = embed(x, self.mask_in, order=self.order)
nrepetitions = x.size // self.nargin
if nrepetitions > 1:
if self.order == "C":
y = xp.zeros((nrepetitions, ) + self.coeff_arr_shape,
dtype=x.dtype)
for rep in range(nrepetitions):
y[rep, ...] = self._forward1(x[rep, ...])
else:
y = xp.zeros(self.coeff_arr_shape + (nrepetitions, ),
dtype=x.dtype)
for rep in range(nrepetitions):
y[..., rep] = self._forward1(x[..., rep])
else:
y = self._forward1(x)
if (self.mask_out is not None) and (not self.nd_output):
y = masker(y, self.mask_out, order=self.order)
return y
def adjoint(self, coeffs):
xp = self.xp
if coeffs.size % self.nargout != 0:
raise ValueError("shape mismatch for adjoint DWT")
nrepetitions = coeffs.size // self.nargout
if (self.mask_out is not None) and (not self.nd_output):
coeffs = embed(coeffs, self.mask_out, order=self.order)
coeffs = coeffs.ravel(order=self.order)
if nrepetitions > 1:
if self.order == "C":
x = xp.zeros((nrepetitions, ) + self.arr_shape,
dtype=coeffs.dtype)
for rep in range(nrepetitions):
x[rep, ...] = self._adjoint1(coeffs[rep, ...])
else:
x = xp.zeros(self.arr_shape + (nrepetitions, ),
dtype=coeffs.dtype)
for rep in range(nrepetitions):
x[..., rep] = self._adjoint1(coeffs[..., rep])
else:
x = self._adjoint1(coeffs)
if (self.mask_in is not None) and (not self.nd_input):
x = masker(x, self.mask_in, order=self.order)
return x
def test_masking_operator(xp, order):
a = xp.ones(100)
mask = xp.ones(a.shape, dtype=xp.bool)
mask[20:40] = 0
# mask_out only
I = MaskingOperator(
mask_out=mask,
nargin=a.size,
squeeze_reps=True,
order=order,
**get_loc(xp),
)
c = I * a
xp.testing.assert_array_equal(masker(a, mask, order=order), c)
r = I.H * c
xp.testing.assert_array_equal(a * mask, r)
# mask_in only
I = MaskingOperator(
mask_in=mask,
nargout=a.size,
squeeze_reps=True,
order=order,
**get_loc(xp),
)
a_masked = masker(a, mask, order=order)
c = I * a_masked
xp.testing.assert_array_equal(a * mask, c)
r = I.H * c
xp.testing.assert_array_equal(a_masked, r)
# mask neither
I = MaskingOperator(nargin=a.size,
squeeze_reps=True,
order=order,
**get_loc(xp))
c = I * a
xp.testing.assert_array_equal(a, c)
r = I.H * c
xp.testing.assert_array_equal(a, r)
# mask both
mask_out = xp.ones(a.shape, dtype=xp.bool)
mask_out[60:80] = 0
# Note: would be identity if same mask used for input & output
I = MaskingOperator(
mask_in=mask,
mask_out=mask_out,
squeeze_reps=True,
order=order,
**get_loc(xp),
)
a_masked = masker(a, mask, order=order)
c = I * a_masked
xp.testing.assert_array_equal(masker(a * mask, mask_out, order=order), c)
r = I.H * c
xp.testing.assert_array_equal(masker(a * mask_out, mask, order=order), r)
# scipy compatibility
xp.testing.assert_array_equal(c, I.matvec(a_masked))
xp.testing.assert_array_equal(r, I.rmatvec(c))
def test_partial_FFT_with_im_mask(xp, nd_in, order, shift):
""" masked FFT with missing samples and masked image domain """
c = get_data(xp)
rstate = xp.random.RandomState(1234)
sample_mask = rstate.rand(*(128, 127)) > 0.5
x, y = xp.meshgrid(
xp.arange(-c.shape[0] // 2, c.shape[0] // 2),
xp.arange(-c.shape[1] // 2, c.shape[1] // 2),
indexing="ij",
sparse=True,
)
# make a circular mask
im_mask = xp.sqrt(x * x + y * y) < c.shape[0] // 2
nd_out = False
FTop = FFT_Operator(
c.shape,
order=order,
im_mask=im_mask,
use_fft_shifts=shift,
nd_input=nd_in,
nd_output=nd_out,
sample_mask=sample_mask,
gpu_force_reinit=False,
mask_kspace_on_gpu=(not shift),
**get_loc(xp),
)
# create new linear operator for forward followed by inverse transform
FtF = FTop.H * FTop
assert isinstance(FtF, LinearOperatorMulti)
# test forward only
forw = embed(FTop * masker(c, im_mask, order=order),
sample_mask,
order=order)
if shift:
expected_forw = sample_mask * fftnc(c * im_mask)
else:
expected_forw = sample_mask * fftn(c * im_mask)
xp.testing.assert_allclose(forw, expected_forw, rtol=1e-7, atol=1e-4)
# test roundtrip
roundtrip = FTop.H * (FTop * masker(c, im_mask, order=order))
if shift:
expected_roundtrip = masker(ifftnc(sample_mask * fftnc(c * im_mask)),
im_mask,
order=order)
else:
expected_roundtrip = masker(ifftn(sample_mask * fftn(c * im_mask)),
im_mask,
order=order)
xp.testing.assert_allclose(roundtrip,
expected_roundtrip,
rtol=1e-7,
atol=1e-4)
# test roundtrip with 2 reps
c2 = xp.stack([c] * 2, axis=-1)
roundtrip = FTop.H * (FTop * masker(c2, im_mask, order=order))
if shift:
expected_roundtrip = masker(
ifftnc(
sample_mask[..., xp.newaxis] *
fftnc(c2 * im_mask[..., xp.newaxis], axes=(0, 1)),
axes=(0, 1),
),
im_mask,
order=order,
)
else:
expected_roundtrip = masker(
ifftn(
sample_mask[..., xp.newaxis] *
fftn(c2 * im_mask[..., xp.newaxis], axes=(0, 1)),
axes=(0, 1),
),
im_mask,
order=order,
)
xp.testing.assert_allclose(roundtrip,
expected_roundtrip,
rtol=1e-7,
atol=1e-4)
def _forward_single_rep(self, x, i_map=0):
xp = self.xp
if self.order == "C":
y_shape = (self.Ncoils, ) + self.shape_in1
else:
y_shape = self.shape_in1 + (self.Ncoils, )
y = xp.zeros(y_shape,
dtype=xp.result_type(x, xp.complex64),
order=self.order)
use_smaps = self.coil_sensitivities is not None
if self.loop_over_coils:
for coil in range(self.Ncoils):
if use_smaps:
if self.order == "C":
xc = x * self.coil_sensitivities[i_map, coil, ...]
else:
xc = x * self.coil_sensitivities[..., coil, i_map]
else:
xc = x
if self.order == "C":
coil_slice = (coil, Ellipsis)
else:
coil_slice = (Ellipsis, coil)
if self.debug:
print("x.shape = {}".format(x.shape))
print("xc.shape = {}".format(xc.shape))
print("y.shape = {}".format(y.shape))
if self.use_fft_shifts:
y[coil_slice] = self.ifftshift(xc, axes=self.fftshift_axes)
else:
y[coil_slice] = xc
if self.preplan_pyfftw:
y[coil_slice] = self.fftn(y[coil_slice])
else:
y[coil_slice] = self.fftn(y[coil_slice],
axes=self.fft_axes)
if self.use_fft_shifts:
y[coil_slice] = self.fftshift(y[coil_slice],
axes=self.fftshift_axes)
if self.basis.transform is not None:
y *= self.basis.transform
else:
if use_smaps:
if self.order == "C":
x = as_complex_array(x)[xp.newaxis, ...]
x = x * self.coil_sensitivities[i_map, ...]
else:
x = as_complex_array(x)[..., xp.newaxis]
x = x * self.coil_sensitivities[..., i_map]
if self.use_fft_shifts:
x = self.ifftshift(x, axes=self.fftshift_axes)
if self.preplan_pyfftw:
y = self.fftn(x)
else:
y = self.fftn(x, axes=self.fft_axes)
if self.use_fft_shifts:
y = self.fftshift(y, axes=self.fftshift_axes)
if self.basis.transform is not None:
if self.order == "C":
y *= self.basis.transform[xp.newaxis, ...]
else:
y *= self.basis.transform[..., xp.newaxis]
if self.sample_mask is not None:
# y = y[self.sample_mask]
y = masker(
y,
mask=self.sample_mask,
order=self.order,
# mask_idx_ravel=self.sample_mask_idx,
)
return y
def test_DWT_masked(xp, order, decimation):
rstate = xp.random.RandomState(5)
c = rstate.randn(128, 128)
# generate a random image-domain mask
# im_mask = xp.random.standard_normal(c.shape) > -0.2
# test 1D in/outputs but using masked values for the input array
im_mask = xp.ones(c.shape, dtype=xp.bool)
im_mask[16:64, 16:64] = 0 # mask out a region
c_masked = masker(c, im_mask, order=order)
fb = filters.pywt_as_filterbank("db2", xp=xp, decimation=decimation)
Wm = MDWT_Operator(
c.shape,
level=2,
filterbank=fb,
order=order,
mode="periodization",
mask_in=im_mask,
mask_out=None,
nd_input=False,
nd_output=False,
random_shift=True,
**get_loc(xp),
)
coeffs = Wm * c_masked
assert_(coeffs.ndim == 1)
if decimation == 2:
assert_(coeffs.size == c.size)
coeffs = coeffs.reshape(c.shape, order=order)
else:
assert_(coeffs.size > c.size)
c_recon = Wm.H * coeffs
assert_(c_recon.ndim == 1)
if xp is np:
assert_(c_recon.size == im_mask.sum())
else:
assert_(c_recon.size == im_mask.sum().get())
c_recon = embed(c_recon, im_mask, order=order)
xp.testing.assert_allclose(c_recon, c * im_mask, rtol=1e-9, atol=1e-9)
# test 1D in/outputs but using masked values for both input and output
# arrays
coeffs_mask = coeffs != 0 # mask out regions of zero-valued coeffs
Wm2 = MDWT_Operator(
c.shape,
level=2,
filterbank=fb,
order=order,
mode="periodization",
mask_in=im_mask,
mask_out=coeffs_mask,
nd_input=False,
nd_output=False,
random_shift=True,
**get_loc(xp),
)
coeffs = Wm2 * c_masked
assert_(coeffs.ndim == 1)
if decimation == 2:
if xp is np:
assert_(coeffs.size == coeffs_mask.sum())
else:
assert_(coeffs.size == coeffs_mask.sum().get())
c_recon = Wm2.H * coeffs
c_recon = embed(c_recon, im_mask, order=order)
xp.testing.assert_allclose(c_recon, c * im_mask, rtol=1e-9, atol=1e-9)