spyrit.core.meas.FreeformSmatrix

class spyrit.core.meas.FreeformSmatrix(meas_shape: int | Size | Iterable[int] = None, M: int = None, index_mask: tensor = None, bool_mask: tensor = None, order: tensor = None, computation: str = 'dyadic', *, noise_model: Module = Identity(), dtype: dtype = torch.float32, device: device = device(type='cpu'))[source]

Bases: FreeformLinear

Simulate linear measurements in a freeform region of interest, using an S-matrix as the acquisition matrix.

\[m =\mathcal{N}\left(Hx\right), \quad \text{where }x = \text{mask}(\tilde{x})\]

where \(\mathcal{N} \colon\, \mathbb{R}^M \to \mathbb{R}^M\) represents a noise operator (e.g., Gaussian), \(S\in\mathbb{R}^{M\times N}\) is a S matrix, \(x \in \mathbb{R}^N\) is the signal in the region of interest, \(M\) is the number of measurements, \(N\) is the number of pixels in the region of interest, \(\text{mask} \colon\, \mathbb{R}^\tilde{N} \to \mathbb{R}^N\) represents the masking operation, \(\tilde{x} \in \mathbb{R}^\tilde{N}\) is the full signal, and \(\tilde{N}\ge N\) is the dimension of the full signal \(\tilde{x}\).

This class plays the same role as FreeformLinear, but instead of accepting an arbitrary measurement matrix \(H\), it sets \(H\) as an S-matrix (see spyrit.misc.walsh_hadamard.walsh_S_matrix()). The S-matrix is built from a Hadamard matrix of order \(N+1\), so \(N+1\) must be a power of two.

Args:

meas_shape (tuple): Shape of the underlying multi-dimensional array \(X\). See FreeformLinear.

M (int, optional): Number of measurements. Defaults to \(N\) (no subsampling), where \(N\) is the number of masked pixels (deduced from index_mask or bool_mask).

index_mask (torch.tensor, optional): See FreeformLinear.

bool_mask (torch.tensor, optional): See FreeformLinear.

order (torch.tensor, optional): Length-\(N\) order vector that defines the measurements to keep (one value per masked pixel). The first component of \(y\) will correspond to the index where order is the highest. Defaults to None (keeps the natural S-matrix row order).

computation (str, optional): Either “dense” or “dyadic”. See the note above for the tradeoff. Defaults to “dyadic”.

noise_model (see spyrit.core.noise): Noise model \(\mathcal{N}\). Defaults to torch.nn.Identity().

dtype (torch.dtype, optional): Data type of the measurement matrix. Defaults to torch.float32.

device (torch.device, optional): Device of the measurement matrix. Defaults to torch.device(“cpu”).

Attributes:

H (torch.tensor): The (subsampled) \(M\times N\) S measurement matrix. When computation is “dense”, this is precomputed and stored; when “dyadic”, it is built on the fly if accessed (e.g. by code expecting the generic Linear interface), which can be memory-heavy for large \(N\) – prefer measure() and fast_pinv(), which never require it.

T (torch.tensor or None): Exact inverse of the full \(N\times N\) S-matrix, used by fast_pinv() when computation is “dense”. None when “dyadic” (not needed: the fast transform is used instead).

computation (str): “dense” or “dyadic”.

M (int): Number of measurements \(M\).

N (int): Number of masked pixels.

order (torch.tensor): Order vector.

indices (torch.tensor): Indices used to reorder the measurement vector.

Note

Choosing computation. Two ways of applying the S-matrix are available, trading off differently depending on \(N\) and the sampling ratio \(M/N\):

  • “dense”: builds and stores the S-matrix explicitly (an \(M\times N\) matrix for H, plus an \(N\times N\) matrix T for the pseudo-inverse), and applies it via a plain matrix-vector product. Cost scales as \(O(MN)\) per measurement/reconstruction. This is the only option available for acquisition matrices that are not built from a power-of-two Hadamard matrix (e.g. a generic, non-dyadic \(H\)); it also becomes memory-heavy for large \(N\) (an \(N\times N\) float32 matrix already exceeds 1 GB around \(N=16000\), and building it can itself fail with an out-of-memory error before any measurement is even taken).

  • “dyadic” (default): uses the fast Walsh-Hadamard-based transform (spyrit.misc.walsh_hadamard.fwalsh_S_torch() / ifwalsh_S_torch()) instead of a matrix-vector product. This requires no \(N\times N\) (or \(M\times N\)) matrix to ever be stored, and costs \(O(N\log N)\) regardless of \(M\) – but that “regardless of \(M\)” is also its main limitation: unlike the dense path, it cannot skip work when subsampling, since it always computes all \(N\) outputs (or requires all \(N\) inputs for the inverse) before the top-\(M\) measurements are selected. It relies on the dyadic (power-of-two) recursive structure of the Hadamard transform, so it is only applicable when \(N+1\) is a power of two – which is always the case for FreeformSmatrix, but would not be for a hypothetical S-matrix-like class built on some other Hadamard matrix whose order is not a power of two.

In practice (see benchmarks in the development notes), the crossover is around \(M/N \approx 0.15\)-0.20, fairly stable across \(N\) from about 1,000 to 16,000: below that sampling ratio, “dense” is faster; above it, “dyadic” is faster (and, for large \(N\), is often the only option that fits in memory at all). If in doubt, benchmark both on your actual \(N\) and \(M\).

Note

As with HadamSmatrix2d, the S-matrix is not orthogonal: the exact inverse used by fast_pinv() is only exact when M equals N (no subsampling); with subsampling, it is an approximation. This holds for both values of computation.

Example 1: Select the first 15 pixels on the diagonal of a batch of images (N`=15, :attr:`N`+1=16=2**4). With full sampling (the default, :attr:`M`=:attr:`N), fast_pinv() exactly recovers the masked pixels, regardless of computation.

>>> h = 32
>>> mask = torch.tensor([[i, i] for i in range(15)]).T
>>> meas_op = FreeformSmatrix(meas_shape=(h, h), index_mask=mask)
>>> print(meas_op.computation)
dyadic
>>> print(meas_op.N, meas_op.M)
15 15
>>> images = torch.rand(4, h, h)
>>> y = meas_op(images)
>>> print(y.shape)
torch.Size([4, 15])
>>> x_hat = meas_op.fast_pinv(y, vectorize=True)
>>> x_true = meas_op.vectorize(images)
>>> print(torch.allclose(x_true, x_hat, atol=1e-4))
True

Example 2: With vectorize = False (the default), the reconstruction is expanded back to the full image shape instead, with unmasked pixels set to fill_value (0 by default).

>>> x_hat_img = meas_op.fast_pinv(y, vectorize=False)
>>> print(x_hat_img.shape)
torch.Size([4, 32, 32])
>>> print(x_hat_img[0, 20, 20].item())  # (20, 20) is not in the mask
0.0

Example 3: With subsampling (M < N), the reconstruction is only approximate (see the note above).

>>> meas_op_sub = FreeformSmatrix(meas_shape=(h, h), M=10, index_mask=mask)
>>> y_sub = meas_op_sub(images)
>>> print(y_sub.shape)
torch.Size([4, 10])
>>> x_hat_sub = meas_op_sub.fast_pinv(y_sub, vectorize=True)
>>> print(torch.allclose(x_true, x_hat_sub, atol=1e-4))
False

Example 4: The two computation modes give the same result (up to floating-point precision), as expected.

>>> meas_op_dense = FreeformSmatrix(meas_shape=(h, h), index_mask=mask, computation="dense")
>>> y_dense = meas_op_dense(images)
>>> print(torch.allclose(y, y_dense, atol=1e-4))
True

Methods

adjoint(m[, unvectorize])

Apply adjoint of matrix H.

fast_pinv(m[, vectorize, fill_value])

Apply the pseudo-inverse of the S measurement matrix.

forward(x)

Simulate measurements.

measure(x)

Simulate noiseless measurements using the S-matrix.

set_matrix_to_inverse(matrix_name)

unvectorize(x[, fill_value])

Unflatten the last dimension of a tensor to the measurement shape at the measured dimensions based on the mask.

vectorize(x)

Appplies the saved mask to the input tensor, where the masked dimensions are collapsed into one.