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:
FreeformLinearSimulate 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 (seespyrit.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\). SeeFreeformLinear.M(int, optional): Number of measurements. Defaults to \(N\) (no subsampling), where \(N\) is the number of masked pixels (deduced fromindex_maskorbool_mask).index_mask(torch.tensor, optional): SeeFreeformLinear.bool_mask(torch.tensor, optional): SeeFreeformLinear.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 whereorderis 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(seespyrit.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. Whencomputationis “dense”, this is precomputed and stored; when “dyadic”, it is built on the fly if accessed (e.g. by code expecting the genericLinearinterface), which can be memory-heavy for large \(N\) – prefermeasure()andfast_pinv(), which never require it.T(torch.tensoror None): Exact inverse of the full \(N\times N\) S-matrix, used byfast_pinv()whencomputationis “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\) matrixTfor 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 forFreeformSmatrix, 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 byfast_pinv()is only exact whenMequalsN(no subsampling); with subsampling, it is an approximation. This holds for both values ofcomputation.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 ofcomputation.>>> 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 tofill_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
computationmodes 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.