From 6bd08850ad92be943a31ac4b096c554571a3f49c Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 1 Nov 2026 21:21:17 +0000 Subject: [PATCH 1/8] CartesianSense: single samples on the grid, with weights positions with one index more than the phase encodes -- (y, x) in 2D, (z, y, x) in 3D -- are single samples on the grid, as GROG-gridded non-Cartesian samples are. A weight per position (e.g. sqrt of a density compensation) multiplies the samples, its conjugate the adjoint, and enters the normal squared; it works for whole readouts too. For single samples the normal kernel is summed straight from the samples, kept only at the places they reach, and applied between FFTs on the image grid (one coset, no doubled grid) through the same cuFFT callbacks, coil kernels and coil slabs as the Cartesian normal. Nothing the size of frames x volume is built; off the card the normal applies the compressed kernel between BART's transforms. Place coordinates become (z, y, x0) in the gather/scatter kernels so lines and points share them. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01MFB6QNUxWSoR8RjtuWacSY --- docs/explanation/non-cartesian.md | 13 + src/bartorch/_abi.py | 1 + src/bartorch/linop/_form.py | 14 +- src/bartorch/linop/_mri.py | 44 +++- src/csrc/include/bartorch.h | 8 +- src/csrc/ops/grid.c | 422 ++++++++++++++++++++++++------ src/csrc/ops/kernels.cu | 42 ++- src/csrc/ops/sense.c | 4 +- tests/test_linop_mri.py | 115 +++++++- 9 files changed, 552 insertions(+), 111 deletions(-) diff --git a/docs/explanation/non-cartesian.md b/docs/explanation/non-cartesian.md index 47edbef64..17030cd4d 100644 --- a/docs/explanation/non-cartesian.md +++ b/docs/explanation/non-cartesian.md @@ -122,6 +122,19 @@ whichever form `--nufft-conf` selects. `toeplitz=False`, or storage forms and the compressed function's mask are in {doc}`../guides/developer/internals/nufft`. +### Samples on the grid + +Samples moved onto the grid beforehand, by GROG-type gridding for example, +need no NUFFT. {func}`~bartorch.linop.CartesianSense` takes them as +`positions` with one index more than the phase encodes, `(z, y, x)` in 3D, +and a weight per sample: its square root of the density compensation, say. +Integer positions make $h$ periodic over the image itself, so the normal is a +per-voxel kernel between FFTs on the image grid, not the doubled one. The +kernel, $\sum_t \lvert w \rvert^2 \overline{\Phi[a', t]}\,\Phi[a, t]$ over +each place's samples, is summed straight from the samples and kept only at the +places they reach; on a card it runs through the same cuFFT callbacks, coil +kernels and coil slabs as the Cartesian normal. + ## Examples - {doc}`../auto_examples/04-non-cartesian/01-trajectories-and-transforms`: diff --git a/src/bartorch/_abi.py b/src/bartorch/_abi.py index 9d0cc3a8c..4c5e4400f 100644 --- a/src/bartorch/_abi.py +++ b/src/bartorch/_abi.py @@ -65,6 +65,7 @@ class Encoding(ctypes.Structure): ("shots", ctypes.c_int64), ("components", ctypes.c_int), ("positions", ctypes.c_void_p), + ("position_weights", ctypes.c_void_p), ("kspace_readout", ctypes.c_int), ("readout", ctypes.c_int64), ("psf", ctypes.c_void_p), diff --git a/src/bartorch/linop/_form.py b/src/bartorch/linop/_form.py index dce21a8fa..bd1e3ebf1 100644 --- a/src/bartorch/linop/_form.py +++ b/src/bartorch/linop/_form.py @@ -198,7 +198,10 @@ class Form: and leaves this as the image's. positions : tensor, optional Phase-encode places as int64 ``(frames, shots, components)``, ``-1`` - for padding, instead of a dense pattern. + for padding, instead of a dense pattern; with one component more than + the phase-encode axes, single samples on the grid. + position_weights : tensor, optional + complex64 ``(frames, shots)`` weights of the positions' entries. coeffs : int Coefficients the image carries, for the plan's report alone. stacked : bool @@ -225,6 +228,7 @@ class Form: stack_positions: torch.Tensor | None = None item_vector: tuple[int, ...] | None = None positions: torch.Tensor | None = None + position_weights: torch.Tensor | None = None frames: int = 1 shots: int = 0 components: int = 0 @@ -306,6 +310,9 @@ def _kspace_factors(self) -> tuple[Factor, ...]: out.append( Factor("positions", tuple(self.positions.shape), ("frames", "samples"), "table") ) + if self.position_weights is not None: + shape = tuple(self.position_weights.shape) + out.append(Factor("position weights", shape, ("frames", "samples"), "table")) if self.weights is not None: out.append(Factor("weights", tuple(self.weights.tensor.shape), ("samples",))) if self.slice_phase is not None: @@ -388,6 +395,8 @@ def array(a: Array | None) -> tuple[int, int]: s.shots = int(self.shots) s.components = int(self.components) s.positions = 0 if self.positions is None else self.positions.data_ptr() + w = self.position_weights + s.position_weights = 0 if w is None else w.data_ptr() s.kspace_readout = int(self.kspace_readout) s.readout = int(self.readout) @@ -435,7 +444,8 @@ def keep(self) -> tuple: self.slice_phase, ) out = [a.tensor for a in arrays if a is not None] - out += [t for t in (self.positions, self.psf, self.stack_positions) if t is not None] + held = (self.positions, self.position_weights, self.psf, self.stack_positions) + out += [t for t in held if t is not None] if self.contraction is not None: out += [self.contraction.sample.tensor, self.contraction.image.tensor] return (*out, *self._keep) diff --git a/src/bartorch/linop/_mri.py b/src/bartorch/linop/_mri.py index 7fb1aaf65..32d26afe2 100644 --- a/src/bartorch/linop/_mri.py +++ b/src/bartorch/linop/_mri.py @@ -199,10 +199,12 @@ def _create(self) -> Built: return build_form(self, self._form()) -def _checked_positions(positions, frames, encodes) -> torch.Tensor: +def _checked_positions(positions, frames, encodes, readout=None) -> torch.Tensor: """Positions as int64 on the host, checked against the frames a basis has and the phase encodes. - Checked before the operator is built, which happens as it is made. + With ``readout``, the readout's length, a position may also carry one + index more, along it: a single sample on the grid. Checked before the + operator is built, which happens as it is made. """ pos = torch.as_tensor(positions) if pos.is_floating_point() or pos.is_complex(): @@ -215,6 +217,8 @@ def _checked_positions(positions, frames, encodes) -> torch.Tensor: raise ValueError(f"positions over frames {lead} need a basis to contract them") if frames is not None and lead != (frames,): raise ValueError(f"positions over frames {lead}, and a basis of {frames}") + if readout is not None and pos.shape[-1] == len(encodes) + 1: + encodes = (*encodes, readout) if pos.shape[-1] != len(encodes): raise ValueError( f"positions of {pos.shape[-1]} indices, " @@ -238,6 +242,13 @@ class _CartesianSampled(_GridSense): nothing the size of the phase-encode plane is held for them. The transforms, the basis and the normal are :class:`_CartesianNative`'s, over the pattern the positions stand for. + + With one index more -- ``(y, x)``, ``(z, y, x)`` -- each position is a + single sample on the grid, as gridded non-Cartesian samples are, and the + samples are ``(*batches, coils, [frames,] shots)``. The normal's kernel + is then summed from the samples and kept only where they fall, and the + normal transforms every spatial axis: one grid, not the doubled one a + NUFFT's Toeplitz normal needs. """ def __init__( @@ -248,6 +259,7 @@ def __init__( basis, readout="kspace", toeplitz=True, + weights=None, **kwargs, ): from bartorch.linop._sense import _grid_ndim @@ -263,7 +275,18 @@ def __init__( sensitivities, image_shape, kwargs.get("kernels", False), kwargs.get("ndim") ) encodes = tuple(image_shape)[len(image_shape) - ndim : -1] - self._positions = _checked_positions(positions, frames, encodes) + self._positions = _checked_positions(positions, frames, encodes, int(image_shape[-1])) + self._points = self._positions.shape[-1] == len(encodes) + 1 + if self._points and readout != "kspace": + raise ValueError("samples on the grid are in k-space along every axis") + self._position_weights = None + if weights is not None: + w = torch.as_tensor(weights) + if tuple(w.shape) != tuple(self._positions.shape[:-1]): + raise ValueError( + f"weights of {tuple(w.shape)}, for positions of {tuple(self._positions.shape)}" + ) + self._position_weights = w.to(device="cpu", dtype=torch.complex64).contiguous() self._kspace_readout = readout == "kspace" self._grid_toeplitz = bool(toeplitz) super().__init__(sensitivities, image_shape, **kwargs) @@ -278,10 +301,12 @@ def _has_basis(self): return self._grid_basis is not None def _kspace_tail(self): - return (*self.encoding, int(self._positions.shape[-2]), int(self.spatial[-1])) + readout = () if self._points else (int(self.spatial[-1]),) + return (*self.encoding, int(self._positions.shape[-2]), *readout) def _kspace_vector(self): - base = {1: self.spatial[-1], 2: int(self._positions.shape[-2]), _layout.COIL: self.coils} + readout = 1 if self._points else self.spatial[-1] + base = {1: readout, 2: int(self._positions.shape[-2]), _layout.COIL: self.coils} return self._encoding_vector(base, self.encoding) def _form(self) -> Form: @@ -295,6 +320,7 @@ def _form(self) -> Form: kernels=self.kernels, basis=None if b is None else Array(b, bvec), positions=pos, + position_weights=self._position_weights, frames=1 if b is None else int(b.shape[1]), shots=int(pos.shape[-2]), components=int(pos.shape[-1]), @@ -851,6 +877,14 @@ def CartesianSense( # noqa: N802 (it is a constructor) encodes, ``([frames,] shots, d)`` with ``d`` indexing ``(y,)`` in 2D or ``(z, y)`` in 3D, and ``-1`` padding frames with fewer shots. The samples are then ``(*batches, [batch,] coils, [frames,] shots, x)``. + With ``d`` one more, indexing ``(y, x)`` or ``(z, y, x)``, each + position is a single sample on the grid -- gridded non-Cartesian + samples, say -- and the samples are ``(*batches, [batch,] coils, + [frames,] shots)``. + weights : tensor, default=None + With ``positions``, a weight per position ``([frames,] shots)``, such + as the square root of a density compensation: it multiplies the + samples, conjugated on the adjoint, and enters the normal squared. readout : {"kspace", "image"}, default='kspace' With ``positions``, whether the samples are in k-space along the readout or already inverse-transformed along it (hybrid space). diff --git a/src/csrc/include/bartorch.h b/src/csrc/include/bartorch.h index 0b29d0a06..335158487 100644 --- a/src/csrc/include/bartorch.h +++ b/src/csrc/include/bartorch.h @@ -315,13 +315,17 @@ struct bartorch_encoding { /* A table of the phase encodes that were sampled, instead of a dense * pattern: `frames` x `shots` places of `components` indices * each -- (y) for a 2D image, (z, y) for a 3D one, -1 for padding -- - * with the whole readout along each. `kspace_readout` says the + * with the whole readout along each; or, with one index more, (y, x) + * or (z, y, x), single samples on the grid. `kspace_readout` says the * samples are in k-space along the readout rather than transformed - * back along it. NULL `positions` is dense samples. */ + * back along it. NULL `positions` is dense samples. + * `position_weights`, `frames` x `shots` complex floats or NULL, + * weight each entry: forward by itself, on the adjoint conjugated. */ int64_t frames; int64_t shots; int components; const int64_t* positions; + const void* position_weights; int kspace_readout; /* A wave: the oversampled readout the coil images are zero-filled to, diff --git a/src/csrc/ops/grid.c b/src/csrc/ops/grid.c index 1dfded18e..472a7ac2e 100644 --- a/src/csrc/ops/grid.c +++ b/src/csrc/ops/grid.c @@ -83,10 +83,12 @@ struct bartorch_grid_axes { extern void bartorch_cuda_bank_to_table(bart_dim_t E, bart_dim_t X, bart_dim_t S, int R, bart_dim_t L, bart_dim_t per, const int* entry_u, const int* u_coord, const unsigned int* mask, const int* prefix, - const struct bartorch_grid_axes* ax, const complex float* B, const complex float* bank, complex float* table); + const struct bartorch_grid_axes* ax, const complex float* B, const complex float* W, + const complex float* bank, complex float* table); extern void bartorch_cuda_table_to_bank(bart_dim_t U, bart_dim_t X, bart_dim_t S, int R, bart_dim_t L, bart_dim_t per, const int* csr_start, const int* csr, const int* u_coord, const unsigned int* mask, const int* prefix, - const struct bartorch_grid_axes* ax, const complex float* B, complex float* bank, const complex float* table); + const struct bartorch_grid_axes* ax, const complex float* B, const complex float* W, + complex float* bank, const complex float* table); #endif static bart_dim_t grid_fused_count = 0; @@ -152,10 +154,14 @@ struct grid_s { /* Sampled-only samples (grid_sampled_create): `E` = `T` frames x `S` * shots, each a phase-encode place with the readout `X` along it, and * `U` distinct places among them. `entry_u` is an entry's place, -1 for - * padding; `u_coord` a place's (z, y); `csr` the entries of each place, - * from `csr_start`. `bt` is the basis as B[t R + r], NULL without one. - * `kspace_readout` says the samples are in k-space along the readout. */ + * padding; `u_coord` a place's (z, y, x0), x0 where its readout starts; + * `csr` the entries of each place, from `csr_start`. `bt` is the basis + * as B[t R + r], NULL without one. `kspace_readout` says the samples + * are in k-space along the readout. With `points` each entry is one + * sample on the grid, `X` is one and x0 its place along the readout. + * `entry_w` weighs each entry, NULL for none. */ bool sampled; + bool points; bool kspace_readout; bart_dim_t T; bart_dim_t S; @@ -167,11 +173,13 @@ struct grid_s { int* csr_start; int* csr; complex float* bt; + complex float* entry_w; int* d_entry_u; int* d_u_coord; int* d_csr_start; int* d_csr; complex float* d_bt; + complex float* d_entry_w; }; static DEF_TYPEID(grid_s); @@ -274,6 +282,8 @@ static complex float* grid_kernel(bart_dim_t kdims[DIMS], const bart_dim_t pat_d * samples, as a bit per place and a count before each word, and the kernel at * those places alone, `kernel[(l R + r) R + c]` -- or no kernel, where every * kept place's is one. */ +static void grid_mods(struct grid_s* d); + static void grid_compress(struct grid_s* d, const bart_dim_t kdims[DIMS], const complex float* K, const bart_dim_t pat_dims[DIMS], const complex float* pattern) { @@ -345,8 +355,13 @@ static void grid_compress(struct grid_s* d, const bart_dim_t kdims[DIMS], const d->kernel = kernel; - /* The centring of each transformed axis, as `fftmod` applies it; ones for - * the uncentred transform. */ + grid_mods(d); +} + +/* The centring of each transformed axis, as `fftmod` applies it; ones for + * the uncentred transform. */ +static void grid_mods(struct grid_s* d) +{ for (int a = 0; a < 3; a++) { d->mod[a] = NULL; @@ -416,7 +431,7 @@ static bool grid_ready(struct grid_s* d, const void* ref) if (d->sampled) { bart_dim_t ed[1] = { MAX(1, d->E) }; - bart_dim_t ud[1] = { MAX(1, 2 * d->U) }; + bart_dim_t ud[1] = { MAX(1, 3 * d->U) }; bart_dim_t sd[1] = { d->U + 1 }; bart_dim_t cd[1] = { MAX(1, d->csr_start[d->U]) }; @@ -436,6 +451,12 @@ static bool grid_ready(struct grid_s* d, const void* ref) d->d_bt = md_alloc_gpu(1, bd, CFL_SIZE); md_copy(1, bd, d->d_bt, d->bt, CFL_SIZE); } + + if (NULL != d->entry_w) { + + d->d_entry_w = md_alloc_gpu(1, ed, CFL_SIZE); + md_copy(1, ed, d->d_entry_w, d->entry_w, CFL_SIZE); + } } d->cb = bartorch_cb_grid_create(d->cim_dims, d->flags, d->kept, @@ -668,11 +689,14 @@ static void sampled_gather(const struct grid_s* d, complex float* out, const com if (0 <= u) { - bart_dim_t at = x * cstr[READ_DIM] + d->u_coord[2 * u + 1] * cstr[PHS1_DIM] - + d->u_coord[2 * u] * cstr[PHS2_DIM] + c * cstr[COIL_DIM]; + bart_dim_t at = (d->u_coord[3 * u + 2] + x) * cstr[READ_DIM] + d->u_coord[3 * u + 1] * cstr[PHS1_DIM] + + d->u_coord[3 * u] * cstr[PHS2_DIM] + c * cstr[COIL_DIM]; for (bart_dim_t r = 0; r < d->R; r++) acc += ((NULL == d->bt) ? 1.f : d->bt[t * d->R + r]) * k[at + r * cstr[COEFF_DIM]]; + + if (NULL != d->entry_w) + acc *= d->entry_w[e]; } out[base + x * ostr[PHS1_DIM]] = acc; @@ -701,8 +725,8 @@ static void sampled_scatter(const struct grid_s* d, complex float* k, const comp for (bart_dim_t x = 0; x < d->X; x++) { - bart_dim_t at = x * cstr[READ_DIM] + d->u_coord[2 * u + 1] * cstr[PHS1_DIM] - + d->u_coord[2 * u] * cstr[PHS2_DIM] + c * cstr[COIL_DIM]; + bart_dim_t at = (d->u_coord[3 * u + 2] + x) * cstr[READ_DIM] + d->u_coord[3 * u + 1] * cstr[PHS1_DIM] + + d->u_coord[3 * u] * cstr[PHS2_DIM] + c * cstr[COIL_DIM]; for (bart_dim_t r = 0; r < d->R; r++) { @@ -715,6 +739,9 @@ static void sampled_scatter(const struct grid_s* d, complex float* k, const comp complex float v = out[x * ostr[PHS1_DIM] + (e - t * d->S) * ostr[PHS2_DIM] + c * ostr[COIL_DIM] + t * ostr[TE_DIM]]; + if (NULL != d->entry_w) + v *= conjf(d->entry_w[e]); + acc += ((NULL == d->bt) ? 1.f : conjf(d->bt[t * d->R + r])) * v; } @@ -819,7 +846,7 @@ static void sampled_fused_forward(struct grid_s* d, complex float* dst, const co grid_in(d, bank + r * per, volume, hybrid, src + c * coil_step + r * coeff_step, m); bartorch_cuda_bank_to_table(d->E, d->X, d->S, (int)d->R, d->kept, per, d->d_entry_u, d->d_u_coord, - d->dmask, d->dprefix, &ax, d->d_bt, bank, table); + d->dmask, d->dprefix, &ax, d->d_bt, d->d_entry_w, bank, table); sampled_readout(d, td, table, true); @@ -876,7 +903,7 @@ static void sampled_fused_adjoint(struct grid_s* d, complex float* dst, const co md_clear(1, bd, bank, CFL_SIZE); bartorch_cuda_table_to_bank(d->U, d->X, d->S, (int)d->R, d->kept, per, d->d_csr_start, d->d_csr, - d->d_u_coord, d->dmask, d->dprefix, &ax, d->d_bt, bank, table); + d->d_u_coord, d->dmask, d->dprefix, &ax, d->d_bt, d->d_entry_w, bank, table); for (bart_dim_t r = 0; r < d->R; r++) grid_out(d, dst + c * coil_step + r * coeff_step, volume, hybrid, bank + r * per, m); @@ -939,6 +966,61 @@ static void grid_adjoint(const linop_data_t* _d, complex float* dst, const compl linop_adjoint(d->fwd, DIMS, d->dom_dims, dst, DIMS, d->out_dims, src); } +/* The normal of samples on the grid off the card: the spectrum of each coil's + * coefficients, at each kept place the kernel across them, nothing at the + * others, and the transform back. */ +static void points_host_normal(const struct grid_s* d, complex float* dst, const complex float* src) +{ + complex float* k = md_alloc(DIMS, d->cim_dims, CFL_SIZE); + md_copy(DIMS, d->cim_dims, k, src, CFL_SIZE); + grid_ft(d, d->cim_dims, d->flags, k, false); + + /* Coils and sets below the coefficients, any batch above them. */ + bart_dim_t R = d->R; + bart_dim_t C = md_calc_size(COEFF_DIM, d->cim_dims) / d->vol; + bart_dim_t O = md_calc_size(DIMS, d->cim_dims) / (d->vol * C * R); + +#pragma omp parallel for collapse(2) + for (bart_dim_t c = 0; c < C * O; c++) { + for (bart_dim_t q = 0; q < d->words; q++) { + + for (bart_dim_t b = 0; (b < 32) && (q * 32 + b < d->vol); b++) { + + bart_dim_t p = q * 32 + b; + complex float* v = k + p + (c % C) * d->vol + (c / C) * d->vol * C * R; + + if (0 == (d->mask[q] & (1u << b))) { + + for (bart_dim_t r = 0; r < R; r++) + v[r * C * d->vol] = 0.f; + + continue; + } + + bart_dim_t l = d->prefix[q] + __builtin_popcount(d->mask[q] & ((1u << b) - 1u)); + complex float in[R]; + + for (bart_dim_t r = 0; r < R; r++) + in[r] = v[r * C * d->vol]; + + for (bart_dim_t r = 0; r < R; r++) { + + complex float acc = 0.f; + + for (bart_dim_t cc = 0; cc < R; cc++) + acc += ((NULL == d->kernel) ? 1.f : d->kernel[(l * R + r) * R + cc]) * in[cc]; + + v[r * C * d->vol] = acc; + } + } + } + } + + grid_ft(d, d->cim_dims, d->flags, k, true); + md_copy(DIMS, d->dom_dims, dst, k, CFL_SIZE); + md_free(k); +} + static void grid_normal(const linop_data_t* _d, complex float* dst, const complex float* src) { const auto d = CAST_DOWN(grid_s, _d); @@ -952,6 +1034,12 @@ static void grid_normal(const linop_data_t* _d, complex float* dst, const comple return; } + if (NULL == d->nrm) { + + points_host_normal(d, dst, src); + return; + } + linop_forward(d->nrm, DIMS, d->dom_dims, dst, DIMS, d->dom_dims, src); } @@ -981,6 +1069,7 @@ static void grid_del(const linop_data_t* _d) md_free(d->d_csr_start); md_free(d->d_csr); md_free(d->d_bt); + md_free(d->d_entry_w); md_free(d->dpsf); xfree(d->entry_u); @@ -988,10 +1077,14 @@ static void grid_del(const linop_data_t* _d) xfree(d->csr_start); xfree(d->csr); xfree(d->bt); + xfree(d->entry_w); md_free(d->psf); - linop_free(d->nrm); - linop_free(d->fwd); + if (NULL != d->nrm) + linop_free(d->nrm); + + if (NULL != d->fwd) + linop_free(d->fwd); xfree(d); } @@ -1257,11 +1350,14 @@ static struct grid_s* grid_state(const bart_dim_t cim_dims[DIMS], d->csr_start = NULL; d->csr = NULL; d->bt = NULL; + d->entry_w = NULL; + d->points = false; d->d_entry_u = NULL; d->d_u_coord = NULL; d->d_csr_start = NULL; d->d_csr = NULL; d->d_bt = NULL; + d->d_entry_w = NULL; debug_printf(DP_DEBUG1, "Cartesian slab: %" PRId64 " coefficients, normal transforms axes %" PRIx64 " of %" PRIx64 ", %" PRId64 " of %" PRId64 " places kept\n", R, d->flags, fft_flags, d->kept, d->plane); @@ -1285,82 +1381,125 @@ const struct linop_s* grid_transform_create(const bart_dim_t cim_dims[DIMS], grid_forward, grid_adjoint, d->toeplitz ? grid_normal : NULL, NULL, grid_del); } -/* What the slab transform over sampled-only k-space holds: `T` frames of `S` shots, each - * a phase-encode position of `components` indices -- (y) for a 2D image, - * (z, y) for a 3D one, -1 for padding -- with the whole readout along it. - * The samples are [1, readout, shots, coils, 1, frames]: in k-space along the - * readout when `kspace_readout`, as the image has it otherwise. +/* What the slab transform over sampled-only k-space holds: `T` frames of `S` + * shots, each a phase-encode position of `components` indices -- (y) for a 2D + * image, (z, y) for a 3D one, -1 for padding -- with the whole readout along + * it; or, with one index more -- (y, x), (z, y, x) -- a single sample on the + * grid. The samples are [1, readout, shots, coils, 1, frames], the readout + * one for single samples: in k-space along the readout when `kspace_readout`, + * as the image has it otherwise. `weights`, one per entry or NULL, multiply + * the samples forward and their conjugates the adjoint. * * The transforms, the basis and the normal are the dense operator's, over the - * pattern the positions stand for: the square root of how often each frame - * samples a place, so a place sampled twice counts twice in the normal as it - * does in the two applications. Forward, the table is read off the gathered - * spectrum and transformed along its readout where that is asked for; that - * is a transform of the table, not of the volume. */ + * pattern the positions stand for: the root of the summed squared weights of + * each frame's entries at a place, so a place sampled twice counts twice in + * the normal as it does in the two applications. Forward, the table is read + * off the gathered spectrum and transformed along its readout where that is + * asked for; that is a transform of the table, not of the volume. + * + * Single samples hold nothing the size of the volume times the frames: the + * kernel is summed straight from the entries of each place the samples + * reach, kept at those places alone, and transformed along every spatial + * axis between its transforms; there is no dense chain, and the normal off + * the card is the two applications. */ static struct grid_s* sampled_state(const bart_dim_t cim_dims[DIMS], bart_dim_t T, bart_dim_t S, int components, - const bart_dim_t* positions, const bart_dim_t bas_dims[DIMS], const complex float* basis, bool centred) + const bart_dim_t* positions, const complex float* weights, + const bart_dim_t bas_dims[DIMS], const complex float* basis, bool centred) { + bart_dim_t nx = cim_dims[READ_DIM]; bart_dim_t ny = cim_dims[PHS1_DIM]; bart_dim_t nz = cim_dims[PHS2_DIM]; + int encodes = (1 < nz) ? 2 : 1; + bool points = (encodes + 1 == components); - if (((1 < nz) ? 2 : 1) != components) - error("bartorch: positions of %d indices for an image of %d phase-encode axes\n", components, (1 < nz) ? 2 : 1); + if ((encodes != components) && !points) + error("bartorch: positions of %d indices for an image of %d phase-encode axes\n", components, encodes); if ((NULL == basis) ? (1 != T) : (bas_dims[TE_DIM] != T)) error("bartorch: positions over %" PRId64 " frames, and a basis of %" PRId64 "\n", T, (NULL == basis) ? 1 : bas_dims[TE_DIM]); - bart_dim_t plane = ny * nz; + /* A place is a readout's (y, z), or a sample's (x, y, z), in the order + * the spectrum lays them out. */ + bart_dim_t places = points ? (nx * ny * nz) : (ny * nz); bart_dim_t E = T * S; - int* place_u = xmalloc((size_t)plane * sizeof(int)); - float* counts = xmalloc((size_t)(T * plane) * sizeof(float)); - int* entry_u = xmalloc((size_t)MAX(1, E) * sizeof(int)); + complex float* w = NULL; - for (bart_dim_t p = 0; p < plane; p++) - place_u[p] = -1; + if (NULL != weights) { - for (bart_dim_t i = 0; i < T * plane; i++) - counts[i] = 0.f; + bart_dim_t ed[1] = { MAX(1, E) }; - bart_dim_t U = 0; - bart_dim_t valid = 0; + w = xmalloc((size_t)ed[0] * sizeof(complex float)); + md_copy(1, ed, w, weights, CFL_SIZE); + } + + int* place_u = xmalloc((size_t)places * sizeof(int)); + int* entry_u = xmalloc((size_t)MAX(1, E) * sizeof(int)); + + for (bart_dim_t p = 0; p < places; p++) + place_u[p] = -1; for (bart_dim_t e = 0; e < E; e++) { - bart_dim_t zc = (2 == components) ? positions[e * components] : 0; - bart_dim_t yc = positions[e * components + components - 1]; + const bart_dim_t* at = positions + e * components; + bool padding = false; - if ((-1 == zc) || (-1 == yc)) { + for (int i = 0; i < components; i++) + padding = padding || (-1 == at[i]); + + if (padding) { entry_u[e] = -1; continue; } - if ((zc < 0) || (zc >= nz) || (yc < 0) || (yc >= ny)) - error("bartorch: a phase encode lies outside the %" PRId64 " x %" PRId64 " plane\n", nz, ny); + bart_dim_t xc = points ? at[components - 1] : 0; + bart_dim_t yc = at[components - 1 - (points ? 1 : 0)]; + bart_dim_t zc = (2 == encodes) ? at[0] : 0; + + if ((zc < 0) || (zc >= nz) || (yc < 0) || (yc >= ny) || (xc < 0) || (xc >= nx)) + error("bartorch: a position lies outside the %" PRId64 " x %" PRId64 " x %" PRId64 " grid\n", nz, ny, nx); - bart_dim_t p = yc + ny * zc; + bart_dim_t p = points ? (xc + nx * (yc + ny * zc)) : (yc + ny * zc); - if (-1 == place_u[p]) + entry_u[e] = (int)p; + place_u[p] = 0; + } + + /* The places in the spectrum's order, so a place's index is where its + * kernel sits among the kept ones. */ + bart_dim_t U = 0; + + for (bart_dim_t p = 0; p < places; p++) + if (0 == place_u[p]) place_u[p] = (int)U++; - entry_u[e] = place_u[p]; - counts[(e / S) * plane + p] += 1.f; + bart_dim_t valid = 0; + + for (bart_dim_t e = 0; e < E; e++) { + + if (0 > entry_u[e]) + continue; + + entry_u[e] = place_u[entry_u[e]]; valid++; } - int* u_coord = xmalloc((size_t)MAX(1, 2 * U) * sizeof(int)); + int* u_coord = xmalloc((size_t)MAX(1, 3 * U) * sizeof(int)); int* csr_start = xmalloc((size_t)(U + 1) * sizeof(int)); int* csr = xmalloc((size_t)MAX(1, valid) * sizeof(int)); int* fill = xmalloc((size_t)MAX(1, U) * sizeof(int)); - for (bart_dim_t p = 0; p < plane; p++) { + for (bart_dim_t p = 0; p < places; p++) { if (0 > place_u[p]) continue; - u_coord[2 * place_u[p]] = (int)(p / ny); - u_coord[2 * place_u[p] + 1] = (int)(p % ny); + int* c = u_coord + 3 * place_u[p]; + + c[0] = (int)(points ? (p / (nx * ny)) : (p / ny)); + c[1] = (int)(points ? ((p / nx) % ny) : (p % ny)); + c[2] = (int)(points ? (p % nx) : 0); } for (bart_dim_t u = 0; u <= U; u++) @@ -1381,55 +1520,162 @@ static struct grid_s* sampled_state(const bart_dim_t cim_dims[DIMS], bart_dim_t csr[fill[entry_u[e]]++] = (int)e; xfree(fill); - xfree(place_u); - bart_dim_t pat_dims[DIMS]; - md_singleton_dims(DIMS, pat_dims); - pat_dims[PHS1_DIM] = ny; - pat_dims[PHS2_DIM] = nz; - pat_dims[TE_DIM] = T; + complex float* bt = NULL; + bart_dim_t R = (NULL == basis) ? 1 : bas_dims[COEFF_DIM]; - complex float* pattern = md_alloc(DIMS, pat_dims, CFL_SIZE); + if (NULL != basis) { - for (bart_dim_t i = 0; i < T * plane; i++) - pattern[i] = sqrtf(counts[i]); + complex float* B = md_alloc(DIMS, bas_dims, CFL_SIZE); + md_copy(DIMS, bas_dims, B, basis, CFL_SIZE); - xfree(counts); + bt = xmalloc((size_t)(T * R) * sizeof(complex float)); - /* Built as for the closed-form normal whatever `toeplitz` says: the - * forward and the adjoint read the kept places it finds. */ - struct grid_s* d = grid_state(cim_dims, pat_dims, pattern, bas_dims, basis, 1, centred); + for (bart_dim_t t = 0; t < T; t++) + for (bart_dim_t r = 0; r < R; r++) + bt[t * R + r] = B[t + T * r]; + + md_free(B); + } + + struct grid_s* d = NULL; + + if (!points) { + + float* counts = xmalloc((size_t)(T * places) * sizeof(float)); + + for (bart_dim_t i = 0; i < T * places; i++) + counts[i] = 0.f; + + for (bart_dim_t u = 0; u < U; u++) { + + bart_dim_t p = u_coord[3 * u + 1] + ny * u_coord[3 * u]; + + for (int q = csr_start[u]; q < csr_start[u + 1]; q++) { + + bart_dim_t e = csr[q]; + float a = (NULL == w) ? 1.f : (crealf(w[e]) * crealf(w[e]) + cimagf(w[e]) * cimagf(w[e])); + + counts[(e / S) * places + p] += a; + } + } + + bart_dim_t pat_dims[DIMS]; + md_singleton_dims(DIMS, pat_dims); + pat_dims[PHS1_DIM] = ny; + pat_dims[PHS2_DIM] = nz; + pat_dims[TE_DIM] = T; + + complex float* pattern = md_alloc(DIMS, pat_dims, CFL_SIZE); + + for (bart_dim_t i = 0; i < T * places; i++) + pattern[i] = sqrtf(counts[i]); + + xfree(counts); + + /* Built as for the closed-form normal whatever `toeplitz` says: + * the forward and the adjoint read the kept places it finds. */ + d = grid_state(cim_dims, pat_dims, pattern, bas_dims, basis, 1, centred); + + md_free(pattern); + + } else { + + bart_flags_t flags = FFT_FLAGS & md_nontriv_dims(DIMS, cim_dims); + + /* cuFFT links no callbacks into a transform along one axis. */ + if (2 > __builtin_popcountl(flags)) + error("bartorch: samples on the grid need an image of two spatial axes or more\n"); + + bart_dim_t pat_dims[DIMS]; + md_singleton_dims(DIMS, pat_dims); + + d = grid_state(cim_dims, pat_dims, NULL, bas_dims, basis, 0, centred); + + linop_free(d->nrm); + d->nrm = NULL; + + d->flags = flags; + d->plane = places; + d->batch = 1; + d->kept = U; + d->words = (places + 31) / 32; + d->mask = xmalloc((size_t)d->words * sizeof(unsigned int)); + d->prefix = xmalloc((size_t)d->words * sizeof(int)); + + bart_dim_t n = 0; + + for (bart_dim_t q = 0; q < d->words; q++) { + + unsigned int bits = 0; + + for (bart_dim_t b = 0; (b < 32) && (q * 32 + b < places); b++) + if (0 <= place_u[q * 32 + b]) + bits |= 1u << b; + + d->mask[q] = bits; + d->prefix[q] = (int)n; + n += __builtin_popcount(bits); + } + + /* kernel[(u R + r) R + c] = sum over the entries e of place u of + * |w_e|^2 conj(B[t, r]) B[t, c], t the frame of e. */ + complex float* kernel = xmalloc((size_t)MAX(1, U * R * R) * sizeof(complex float)); + bool ones = (1 == R); + +#pragma omp parallel for reduction(&&:ones) + for (bart_dim_t u = 0; u < U; u++) { + + for (bart_dim_t r = 0; r < R; r++) { + for (bart_dim_t c = 0; c < R; c++) { + + complex double acc = 0.; + + for (int q = csr_start[u]; q < csr_start[u + 1]; q++) { + + bart_dim_t e = csr[q]; + bart_dim_t t = e / S; + double a = (NULL == w) ? 1. : (double)(crealf(w[e]) * crealf(w[e]) + cimagf(w[e]) * cimagf(w[e])); + + acc += a * ((NULL == bt) ? 1.f : (conjf(bt[t * R + r]) * bt[t * R + c])); + } + + kernel[(u * R + r) * R + c] = (complex float)acc; + ones = ones && (1. == acc); + } + } + } + + if (ones) { + + xfree(kernel); + kernel = NULL; + } + + d->kernel = kernel; + + grid_mods(d); + } + + xfree(place_u); /* The dense forward is not what this operator applies. */ linop_free(d->fwd); d->fwd = NULL; - md_free(pattern); - d->sampled = true; + d->points = points; d->T = T; d->S = S; - d->X = cim_dims[READ_DIM]; + d->X = points ? 1 : nx; d->E = E; d->U = U; d->entry_u = entry_u; d->u_coord = u_coord; d->csr_start = csr_start; d->csr = csr; - - if (NULL != basis) { - - complex float* B = md_alloc(DIMS, bas_dims, CFL_SIZE); - md_copy(DIMS, bas_dims, B, basis, CFL_SIZE); - - d->bt = xmalloc((size_t)(T * d->R) * sizeof(complex float)); - - for (bart_dim_t t = 0; t < T; t++) - for (bart_dim_t r = 0; r < d->R; r++) - d->bt[t * d->R + r] = B[t + T * r]; - - md_free(B); - } + d->bt = bt; + d->entry_w = w; md_select_dims(DIMS, COIL_FLAG, d->out_dims, cim_dims); d->out_dims[PHS1_DIM] = d->X; @@ -1443,10 +1689,13 @@ static struct grid_s* sampled_state(const bart_dim_t cim_dims[DIMS], bart_dim_t } const struct linop_s* grid_sampled_create(const bart_dim_t cim_dims[DIMS], bart_dim_t T, bart_dim_t S, int components, - const bart_dim_t* positions, const bart_dim_t bas_dims[DIMS], const complex float* basis, + const bart_dim_t* positions, const complex float* weights, const bart_dim_t bas_dims[DIMS], const complex float* basis, int kspace_readout, int toeplitz) { - struct grid_s* d = sampled_state(cim_dims, T, S, components, positions, bas_dims, basis, true); + struct grid_s* d = sampled_state(cim_dims, T, S, components, positions, weights, bas_dims, basis, true); + + if (d->points && (0 == kspace_readout)) + error("bartorch: samples on the grid are in k-space along every axis\n"); d->toeplitz = (0 != toeplitz); d->kspace_readout = (0 != kspace_readout); @@ -1551,7 +1800,10 @@ const struct linop_s* wave_sampled_create(const bart_dim_t dom_dims[DIMS], bart_ md_copy_dims(DIMS, cim_dims, dom_dims); cim_dims[READ_DIM] = wx; - struct grid_s* d = sampled_state(cim_dims, T, S, components, positions, bas_dims, basis, 0 != centred); + struct grid_s* d = sampled_state(cim_dims, T, S, components, positions, NULL, bas_dims, basis, 0 != centred); + + if (d->points) + error("bartorch: a wave samples whole readouts\n"); wave_setup(d, dom_dims, psf); diff --git a/src/csrc/ops/kernels.cu b/src/csrc/ops/kernels.cu index 7506fced1..09d777aa9 100644 --- a/src/csrc/ops/kernels.cu +++ b/src/csrc/ops/kernels.cu @@ -358,11 +358,13 @@ __device__ static inline bart_dim_t grid_spectrum_at(const struct bartorch_grid_ return (0 > j) ? -1 : batch * L + j; } -/* table[e X + x] = centring * sum_r B[t R + r] bank[r per + at(x, place of e)], - * with t = e / S, zero for a padding entry. B NULL is one coefficient. */ +/* table[e X + x] = W[e] * centring * sum_r B[t R + r] bank[r per + at(x0 + x, place of e)], + * with t = e / S, zero for a padding entry. B NULL is one coefficient, W + * NULL no weights; a place's coordinates are (z, y, x0). */ __global__ static void kern_bank_to_table(bart_dim_t E, bart_dim_t X, bart_dim_t S, int R, bart_dim_t L, bart_dim_t per, const int* entry_u, const int* u_coord, const unsigned int* mask, const int* prefix, - struct bartorch_grid_axes ax, const cuFloatComplex* B, const cuFloatComplex* bank, cuFloatComplex* table) + struct bartorch_grid_axes ax, const cuFloatComplex* B, const cuFloatComplex* W, + const cuFloatComplex* bank, cuFloatComplex* table) { bart_dim_t start = threadIdx.x + (bart_dim_t)blockDim.x * blockIdx.x; bart_stride_t stride = (bart_dim_t)blockDim.x * gridDim.x; @@ -377,11 +379,14 @@ __global__ static void kern_bank_to_table(bart_dim_t E, bart_dim_t X, bart_dim_t if (0 <= u) { - bart_dim_t c[3] = { i - e * X, u_coord[2 * u + 1], u_coord[2 * u] }; + bart_dim_t c[3] = { u_coord[3 * u + 2] + i - e * X, u_coord[3 * u + 1], u_coord[3 * u] }; cuFloatComplex m; bart_dim_t at = grid_spectrum_at(&ax, L, mask, prefix, c, &m); + if (NULL != W) + m = cuCmulf(m, W[e]); + if (0 <= at) { bart_dim_t t = e / S; @@ -401,12 +406,13 @@ __global__ static void kern_bank_to_table(bart_dim_t E, bart_dim_t X, bart_dim_t } } -/* The adjoint: bank[r per + at(x, u)] = conj(centring) * sum over the entries - * e of place u of conj(B[t R + r]) table[e X + x]. Each (u, x) is its own - * place of the spectrum, so no two threads write the same one. */ +/* The adjoint: bank[r per + at(x0 + x, u)] = conj(centring) * sum over the + * entries e of place u of conj(W[e] B[t R + r]) table[e X + x]. Each (u, x) + * is its own place of the spectrum, so no two threads write the same one. */ __global__ static void kern_table_to_bank(bart_dim_t U, bart_dim_t X, bart_dim_t S, int R, bart_dim_t L, bart_dim_t per, const int* csr_start, const int* csr, const int* u_coord, const unsigned int* mask, const int* prefix, - struct bartorch_grid_axes ax, const cuFloatComplex* B, cuFloatComplex* bank, const cuFloatComplex* table) + struct bartorch_grid_axes ax, const cuFloatComplex* B, const cuFloatComplex* W, + cuFloatComplex* bank, const cuFloatComplex* table) { bart_dim_t start = threadIdx.x + (bart_dim_t)blockDim.x * blockIdx.x; bart_stride_t stride = (bart_dim_t)blockDim.x * gridDim.x; @@ -415,7 +421,8 @@ __global__ static void kern_table_to_bank(bart_dim_t U, bart_dim_t X, bart_dim_t for (bart_dim_t i = start; i < n; i += stride) { bart_dim_t u = i / X; - bart_dim_t c[3] = { i - u * X, u_coord[2 * u + 1], u_coord[2 * u] }; + bart_dim_t x = i - u * X; + bart_dim_t c[3] = { u_coord[3 * u + 2] + x, u_coord[3 * u + 1], u_coord[3 * u] }; cuFloatComplex m; bart_dim_t at = grid_spectrum_at(&ax, L, mask, prefix, c, &m); @@ -432,7 +439,10 @@ __global__ static void kern_table_to_bank(bart_dim_t U, bart_dim_t X, bart_dim_t for (bart_dim_t k = csr_start[u]; k < csr_start[u + 1]; k++) { bart_dim_t e = csr[k]; - cuFloatComplex v = table[e * X + c[0]]; + cuFloatComplex v = table[e * X + x]; + + if (NULL != W) + v = cuCmulf(cuConjf(W[e]), v); acc = cuCaddf(acc, (NULL == B) ? v : cuCmulf(cuConjf(B[(e / S) * R + r]), v)); } @@ -444,20 +454,24 @@ __global__ static void kern_table_to_bank(bart_dim_t U, bart_dim_t X, bart_dim_t extern "C" void bartorch_cuda_bank_to_table(bart_dim_t E, bart_dim_t X, bart_dim_t S, int R, bart_dim_t L, bart_dim_t per, const int* entry_u, const int* u_coord, const unsigned int* mask, const int* prefix, - const struct bartorch_grid_axes* ax, const _Complex float* B, const _Complex float* bank, _Complex float* table) + const struct bartorch_grid_axes* ax, const _Complex float* B, const _Complex float* W, + const _Complex float* bank, _Complex float* table) { kern_bank_to_table<<>>(E, X, S, R, L, per, entry_u, u_coord, - mask, prefix, *ax, (const cuFloatComplex*)B, (const cuFloatComplex*)bank, (cuFloatComplex*)table); + mask, prefix, *ax, (const cuFloatComplex*)B, (const cuFloatComplex*)W, + (const cuFloatComplex*)bank, (cuFloatComplex*)table); CUDA_KERNEL_ERROR; } extern "C" void bartorch_cuda_table_to_bank(bart_dim_t U, bart_dim_t X, bart_dim_t S, int R, bart_dim_t L, bart_dim_t per, const int* csr_start, const int* csr, const int* u_coord, const unsigned int* mask, const int* prefix, - const struct bartorch_grid_axes* ax, const _Complex float* B, _Complex float* bank, const _Complex float* table) + const struct bartorch_grid_axes* ax, const _Complex float* B, const _Complex float* W, + _Complex float* bank, const _Complex float* table) { kern_table_to_bank<<>>(U, X, S, R, L, per, csr_start, csr, u_coord, - mask, prefix, *ax, (const cuFloatComplex*)B, (cuFloatComplex*)bank, (const cuFloatComplex*)table); + mask, prefix, *ax, (const cuFloatComplex*)B, (const cuFloatComplex*)W, + (cuFloatComplex*)bank, (const cuFloatComplex*)table); CUDA_KERNEL_ERROR; } diff --git a/src/csrc/ops/sense.c b/src/csrc/ops/sense.c index 7e4edd3bf..490466ca4 100644 --- a/src/csrc/ops/sense.c +++ b/src/csrc/ops/sense.c @@ -1522,7 +1522,7 @@ extern const struct linop_s* grid_transform_create(const bart_dim_t cim_dims[DIM const bart_dim_t pat_dims[DIMS], const complex float* pattern, const bart_dim_t bas_dims[DIMS], const complex float* basis, int toeplitz); extern const struct linop_s* grid_sampled_create(const bart_dim_t cim_dims[DIMS], bart_dim_t T, bart_dim_t S, int components, - const bart_dim_t* positions, const bart_dim_t bas_dims[DIMS], const complex float* basis, + const bart_dim_t* positions, const complex float* weights, const bart_dim_t bas_dims[DIMS], const complex float* basis, int kspace_readout, int toeplitz); extern const struct linop_s* wave_transform_create(const bart_dim_t dom_dims[DIMS], bart_dim_t wx, const complex float* psf, int centred, const bart_dim_t pat_dims[DIMS], const complex float* pattern, @@ -1655,7 +1655,7 @@ static const struct linop_s* form_transform(const struct bartorch_encoding* f, c if (NULL != f->positions) return grid_sampled_create(cim_dims, f->frames, f->shots, f->components, f->positions, - f->bas_dims, f->basis, f->kspace_readout, f->toeplitz); + f->position_weights, f->bas_dims, f->basis, f->kspace_readout, f->toeplitz); /* With no k-space factor the transform is the whole of it: * centred, which is what `bartorch.tools.fft` is and so what a diff --git a/tests/test_linop_mri.py b/tests/test_linop_mri.py index b316a7645..5b58cb77a 100644 --- a/tests/test_linop_mri.py +++ b/tests/test_linop_mri.py @@ -834,7 +834,7 @@ def test_sampled_positions_are_checked(maps, positions_2d): with pytest.raises(ValueError, match="pattern or positions"): linop.CartesianSense(maps, (Y, X), pattern=torch.ones(Y, 1), positions=positions_2d) with pytest.raises(ValueError, match="indices"): - linop.CartesianSense(maps, (Y, X), positions=torch.zeros(4, 2, dtype=torch.long)) + linop.CartesianSense(maps, (Y, X), positions=torch.zeros(4, 3, dtype=torch.long)) with pytest.raises(ValueError, match="outside"): linop.CartesianSense(maps, (Y, X), positions=torch.full((4, 1), Y)) with pytest.raises(ValueError, match="basis"): @@ -895,6 +895,119 @@ def test_on_a_card_a_2d_sampled_encoding_is_the_host_one(maps, positions_2d): assert (A.adjoint(samples) - want).abs().max() / want.abs().max() < 1e-5 +# --- single samples on the grid --------------------------------------------- +# +# A position with one index more than the phase encodes -- (z, y, x) in 3D -- +# is one sample on the grid, as gridded non-Cartesian samples are; a weight +# per position multiplies it and enters the normal squared. + + +def _points(frames, shots, shape, seed): + """Positions (frames, shots, 3), some repeated, the second frame padded, and a weight each.""" + torch.manual_seed(seed) + out = torch.stack([torch.randint(0, n, (frames, shots)) for n in shape], -1) + out[1, -2:] = -1 + weights = (torch.rand(frames, shots) + 0.5).to(torch.complex64) + return out, weights + + +def _read_points(spectra, positions, weights, b): + """(coils, coeffs, z, y, x) spectra read at the positions, contracted with b.""" + frames, shots = positions.shape[:2] + out = torch.zeros(spectra.shape[0], frames, shots, dtype=spectra.dtype) + for t in range(frames): + for s in range(shots): + z, y, x = (int(i) for i in positions[t, s]) + if z >= 0: + out[:, t, s] = weights[t, s] * (b[:, t] * spectra[:, :, z, y, x]).sum(1) + return out + + +def _points_pattern(positions, weights, shape): + """The dense pattern with the same normal: per frame, the root of the summed squared weights.""" + frames, shots = positions.shape[:2] + out = torch.zeros(frames, *shape, dtype=torch.float64) + for t in range(frames): + for s in range(shots): + index = tuple(int(i) for i in positions[t, s]) + if min(index) >= 0: + out[(t, *index)] += float(weights[t, s].abs()) ** 2 + return out.sqrt().to(torch.complex64) + + +def test_points_are_the_dense_spectrum_at_the_positions_and_the_normal_is_the_dense_one(): + coils, z, y, x, frames, coeffs, shots = 3, 4, 6, 5, 5, 2, 30 + torch.manual_seed(21) + maps, b, image = _rand(coils, z, y, x), _rand(coeffs, frames), _rand(coeffs, z, y, x) + positions, weights = _points(frames, shots, (z, y, x), 22) + A = linop.CartesianSense(maps, (coeffs, z, y, x), positions=positions, weights=weights, basis=b) + assert A.oshape == (coils, frames, shots) + + want = _read_points(_fftc(maps[:, None] * image, (-3, -2, -1)), positions, weights, b) + got = A(image) + assert (got - want).abs().max() / want.abs().max() < 1e-5 + assert _adjointness(A) < 1e-5 + + slow = linop.CartesianSense( + maps, (coeffs, z, y, x), positions=positions, weights=weights, basis=b, toeplitz=False + ) + dense = linop.CartesianSense( + maps, (coeffs, z, y, x), pattern=_points_pattern(positions, weights, (z, y, x)), basis=b + ) + want = slow.adjoint(slow(image)) + for got in (A.normal(image), dense.normal(image)): + assert (got - want).abs().max() / want.abs().max() < 1e-5 + + +def test_points_in_2d_without_a_basis_or_weights(maps): + torch.manual_seed(23) + positions = torch.stack([torch.randint(0, Y, (12,)), torch.randint(0, X, (12,))], -1) + A = linop.CartesianSense(maps, (Y, X), positions=positions) + assert A.oshape == (COILS, 12) + image = _rand(Y, X) + want = A.adjoint(A(image)) + assert (A.normal(image) - want).abs().max() / want.abs().max() < 1e-5 + assert _adjointness(A) < 1e-5 + + +def test_points_are_checked(maps): + with pytest.raises(ValueError, match="k-space along every axis"): + linop.CartesianSense( + maps, (Y, X), positions=torch.zeros(4, 2, dtype=torch.long), readout="image" + ) + with pytest.raises(ValueError, match="weights of"): + linop.CartesianSense( + maps, (Y, X), positions=torch.zeros(4, 2, dtype=torch.long), weights=torch.ones(3) + ) + + +@requires_cuda +def test_on_a_card_points_run_through_cufft_and_are_the_host_ones(): + coils, z, y, x, frames, coeffs, shots = 3, 4, 6, 5, 5, 2, 30 + torch.manual_seed(24) + maps, b, image = _rand(coils, z, y, x), _rand(coeffs, frames), _rand(coeffs, z, y, x) + positions, weights = _points(frames, shots, (z, y, x), 25) + kwargs = dict(positions=positions, weights=weights, basis=b) + A = linop.CartesianSense(maps, (coeffs, z, y, x), device="cuda", **kwargs) + host = linop.CartesianSense(maps, (coeffs, z, y, x), **kwargs) + + before = library().bartorch_grid_fused() + got = A(image) + assert library().bartorch_grid_fused() > before, "the forward ran through cuFFT's callbacks" + want = host(image) + assert (got - want).abs().max() / want.abs().max() < 1e-5 + + samples = _rand(*A.oshape) + want = host.adjoint(samples) + assert (A.adjoint(samples) - want).abs().max() / want.abs().max() < 1e-5 + + before = library().bartorch_grid_fused() + got = A.normal(image) + assert library().bartorch_grid_fused() > before, "the normal ran through cuFFT's callbacks" + want = host.normal(image) + assert (got - want).abs().max() / want.abs().max() < 1e-5 + + # --- wave in the coil loop --------------------------------------------------- # # The table of a wave encoding is the dense one read at the positions, and its From bfd8cf2858d2876a7ad1da78d3a072c214831bdf Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 1 Nov 2026 22:21:42 +0000 Subject: [PATCH 2/8] data_scaling: find BART's order statistics by selection, on the card With an operator, data_scaling handed |A^H y| to BART's estimate_scaling_norm, which copies it to the host and sorts the whole array for three order statistics. At 256^3 x rank 4 that is 34 s of a recon. The median, the 90th percentile, the maximum (or the asked-for percentile) are now found with torch.kthvalue at BART's indices, on the operator's card when it has one: the same rule and the same elements, to the last bit of the magnitude. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01MFB6QNUxWSoR8RjtuWacSY --- src/bartorch/optim/_scaling.py | 52 +++++++++++++++++++++++++--------- tests/test_solve.py | 26 +++++++++++++++++ 2 files changed, 65 insertions(+), 13 deletions(-) diff --git a/src/bartorch/optim/_scaling.py b/src/bartorch/optim/_scaling.py index 088639402..be07d3e65 100644 --- a/src/bartorch/optim/_scaling.py +++ b/src/bartorch/optim/_scaling.py @@ -5,8 +5,6 @@ import torch from bartorch import _call -from bartorch._dispatch import _ensure_ready, _lock -from bartorch._lib import library from bartorch._operator import as_operand _estscaling = _call.build("estscaling", __name__) @@ -50,6 +48,10 @@ def data_scaling( Notes ----- + The order statistics are found by selection (``torch.kthvalue``), on the + operator's card when it has one, rather than by BART's sort of the whole + image on the host; they are the same values. + Without ``A`` the estimate is BART's ``estimate_scaling``: the root-sum-of-squares image of the fully sampled central region of k-space, at most 32 samples along each axis, corrected for the region's size, and the @@ -74,14 +76,38 @@ def data_scaling( adjoint = A.adjoint(y) image = as_operand(adjoint, tuple(adjoint.shape), "the adjoint image") - _ensure_ready() - with _lock: - return float( - library().bartorch_scaling_norm( - image.numel(), - image.data_ptr(), - 1.0, - int(compat), - -1.0 if percentile is None else float(percentile), - ) - ) + magnitudes = image.reshape(-1).abs() + card = getattr(A, "device", None) + if card is not None and torch.device(card).type == "cuda" and magnitudes.device.type == "cpu": + # Selection over tens of millions of values is seconds on the host + # and milliseconds on the card the operator already uses. + magnitudes = magnitudes.to(card) + return _scaling_norm(magnitudes, compat, percentile) + + +def _scaling_norm(magnitudes: torch.Tensor, compat: bool, percentile: float | None) -> float: + """BART's ``estimate_scaling_norm`` over ``magnitudes``, by selection where they lie. + + The same order statistics BART sorts the whole array for -- the median at + ``n // 2``, the 90th percentile at ``trunc(0.9 n)``, the maximum, or the + asked-for percentile -- found with ``kthvalue``, on the card when the + image is there. + """ + n = magnitudes.numel() + if n == 0: + return 0.0 + m = magnitudes.to(torch.float32) + + def at(i: int) -> float: + return float(torch.kthvalue(m, min(n - 1, max(0, i)) + 1).values) + + if percentile is not None and percentile > 0: + if percentile > 1: + raise ValueError(f"the percentile is a fraction, at most one, not {percentile}") + return at(int(n * percentile)) + median = at(n // 2) + if compat: + return median + p90 = at(int(n * 0.9)) + top = float(m.max()) + return p90 if (top - p90) < 2 * (p90 - median) else top diff --git a/tests/test_solve.py b/tests/test_solve.py index 019d8777d..48cb2ff49 100644 --- a/tests/test_solve.py +++ b/tests/test_solve.py @@ -764,6 +764,32 @@ def test_the_scaling_for_a_trajectory_is_the_other_branch(): assert optim.data_scaling(y * 4.0, A=A) == pytest.approx(4 * scale, rel=1e-5) +@pytest.mark.parametrize("compat, percentile", [(False, None), (True, None), (False, 0.99)]) +def test_the_scaling_by_selection_is_bart_sorting_the_magnitudes(compat, percentile): + """data_scaling with an operator finds BART's order statistics by selection; + they are those BART's estimate_scaling_norm sorts the image for, to the + last bit of the magnitude (torch's and the C library's ``cabsf`` differ there).""" + from bartorch._dispatch import _ensure_ready + from bartorch._lib import library + + n = 16 + traj = bt.traj(readout=n, spokes=24, radial=True) + maps = _rand(2, n, n) + A = linop.NoncartesianSense(maps, (n, n), traj=traj) + y = A(_rand(n, n)) + image = A.adjoint(y).contiguous() + _ensure_ready() + want = library().bartorch_scaling_norm( + image.numel(), + image.data_ptr(), + 1.0, + int(compat), + -1.0 if percentile is None else percentile, + ) + got = optim.data_scaling(y, A=A, compat=compat, percentile=percentile) + assert got == pytest.approx(want, rel=1e-6) + + def test_a_wavelet_term_reused_answers_as_a_freshly_built_one(): """BART's wavelet threshold draws its shifts from a generator of its own; the solve rewinds it, so a kept term answers as a fresh one would.""" From 1ddc635ab13cbbdf92c19f3fedde2a0a4ae98d4f Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 1 Nov 2026 22:32:34 +0000 Subject: [PATCH 3/8] test: data_scaling of an empty image and a percentile above one Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01MFB6QNUxWSoR8RjtuWacSY --- tests/test_solve.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tests/test_solve.py b/tests/test_solve.py index 48cb2ff49..79cabd93f 100644 --- a/tests/test_solve.py +++ b/tests/test_solve.py @@ -790,6 +790,14 @@ def test_the_scaling_by_selection_is_bart_sorting_the_magnitudes(compat, percent assert got == pytest.approx(want, rel=1e-6) +def test_the_scaling_of_nothing_is_zero_and_a_percentile_is_a_fraction(): + from bartorch.optim._scaling import _scaling_norm + + assert _scaling_norm(torch.empty(0), False, None) == 0.0 + with pytest.raises(ValueError, match="fraction"): + _scaling_norm(torch.ones(4), False, 1.5) + + def test_a_wavelet_term_reused_answers_as_a_freshly_built_one(): """BART's wavelet threshold draws its shifts from a generator of its own; the solve rewinds it, so a kept term answers as a fresh one would.""" From 49f89612092930bc2ce66325fa873cbbfc8a1222 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 1 Nov 2026 23:00:57 +0000 Subject: [PATCH 4/8] optim: data_scaling takes a given A^H y; mean_diagonal draws and applies its probes on a device A solve's start computes A^H y anyway, and on a card a volume's probe crossing the bus twice cost a second of a 256^3 recon. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01MFB6QNUxWSoR8RjtuWacSY --- src/bartorch/optim/_linear.py | 14 ++++++++++---- src/bartorch/optim/_scaling.py | 7 ++++++- tests/test_solve.py | 20 ++++++++++++++++++++ 3 files changed, 36 insertions(+), 5 deletions(-) diff --git a/src/bartorch/optim/_linear.py b/src/bartorch/optim/_linear.py index ed92c92e1..6986f2bda 100644 --- a/src/bartorch/optim/_linear.py +++ b/src/bartorch/optim/_linear.py @@ -232,7 +232,7 @@ def maxeigen( return float(value.value) -def mean_diagonal(A, *, probes: int = 1, seed: int = 0) -> float: +def mean_diagonal(A, *, probes: int = 1, seed: int = 0, device=None) -> float: r"""Estimate the mean of the diagonal of :math:`A^H A` from random probes (Hutchinson). For a probe :math:`v` of independent standard complex normal entries, @@ -252,7 +252,12 @@ def mean_diagonal(A, *, probes: int = 1, seed: int = 0) -> float: :math:`1/\sqrt{n}` for an operator of :math:`n` unknowns with a flat diagonal, so a volume needs one. seed : int, default=0 - Seed of the host generator the probes are drawn from. + Seed of the generator the probes are drawn from. + device : torch.device or str, default=None + Where the probes are drawn, applied and reduced; on the host when + omitted. On the operator's card a volume's probe never crosses the + bus, and the generator is the card's, so the estimate differs from the + host's by the probe's spread. Returns ------- @@ -265,10 +270,11 @@ def mean_diagonal(A, *, probes: int = 1, seed: int = 0) -> float: """ if probes < 1: raise ValueError(f"the estimate takes at least one probe, not {probes}") - generator = torch.Generator().manual_seed(int(seed)) + on = torch.device("cpu" if device is None else device) + generator = torch.Generator(device=on).manual_seed(int(seed)) total = 0.0 for _ in range(int(probes)): - v = torch.randn(A.ishape, dtype=torch.complex64, generator=generator) + v = torch.randn(A.ishape, dtype=torch.complex64, generator=generator, device=on) total += float(torch.vdot(v.reshape(-1), A.normal(v).reshape(-1)).real / v.norm() ** 2) return total / probes diff --git a/src/bartorch/optim/_scaling.py b/src/bartorch/optim/_scaling.py index be07d3e65..baac97989 100644 --- a/src/bartorch/optim/_scaling.py +++ b/src/bartorch/optim/_scaling.py @@ -18,6 +18,7 @@ def data_scaling( A=None, percentile: float | None = None, compat: bool = False, + adjoint: torch.Tensor | None = None, ) -> float: """Estimate the factor by which BART's reconstructions divide the data. @@ -39,6 +40,9 @@ def data_scaling( Percentile of the sorted magnitudes to take in place of BART's rule. compat : bool, default=False Take the median, BART's legacy estimate. Only with ``A``. + adjoint : torch.Tensor, default=None + ``A^H y`` when the caller has it already, as a solve's start does; + ``A`` is then not applied again. Only with ``A``. Returns ------- @@ -74,7 +78,8 @@ def data_scaling( estimate = _estscaling(y, **({} if percentile is None else {"percentile": percentile})) return float(estimate.reshape(-1)[0].abs()) - adjoint = A.adjoint(y) + if adjoint is None: + adjoint = A.adjoint(y) image = as_operand(adjoint, tuple(adjoint.shape), "the adjoint image") magnitudes = image.reshape(-1).abs() card = getattr(A, "device", None) diff --git a/tests/test_solve.py b/tests/test_solve.py index 79cabd93f..e7e029710 100644 --- a/tests/test_solve.py +++ b/tests/test_solve.py @@ -790,6 +790,17 @@ def test_the_scaling_by_selection_is_bart_sorting_the_magnitudes(compat, percent assert got == pytest.approx(want, rel=1e-6) +def test_a_given_adjoint_is_the_scaling_of_the_data_it_came_from(): + n = 16 + traj = bt.traj(readout=n, spokes=24, radial=True) + A = linop.NoncartesianSense(_rand(2, n, n), (n, n), traj=traj) + y = A(_rand(n, n)) + want = optim.data_scaling(y, A=A) + given = A.adjoint(y) + assert optim.data_scaling(y, A=A, adjoint=given) == pytest.approx(want, rel=1e-6) + assert optim.data_scaling(None, A=A, adjoint=2 * given) == pytest.approx(2 * want, rel=1e-6) + + def test_the_scaling_of_nothing_is_zero_and_a_percentile_is_a_fraction(): from bartorch.optim._scaling import _scaling_norm @@ -1045,10 +1056,19 @@ def test_the_mean_diagonal_is_the_trace_over_the_size(_small_encoding): exact = float(torch.diagonal(M.conj().T @ M).real.mean()) assert abs(optim.mean_diagonal(A, probes=64) - exact) / exact < 0.05 assert optim.mean_diagonal(A, seed=3) == optim.mean_diagonal(A, seed=3) + assert optim.mean_diagonal(A, seed=3, device="cpu") == optim.mean_diagonal(A, seed=3) with pytest.raises(ValueError, match="at least one probe"): optim.mean_diagonal(A, probes=0) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a card") +def test_the_mean_diagonal_on_the_card_is_the_trace_over_the_size(_small_encoding): + A, _ = _small_encoding + M = _dense(A) + exact = float(torch.diagonal(M.conj().T @ M).real.mean()) + assert abs(optim.mean_diagonal(A, probes=64, device="cuda") - exact) / exact < 0.05 + + def test_a_solve_placed_elsewhere_moves_its_operands_and_returns_home(): """What ``device=`` does around a solve, observed through the operands it is handed.""" seen = [] From 94d38544fe008c034e0f73ef83aae633fe7c15d2 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 1 Nov 2026 23:19:02 +0000 Subject: [PATCH 5/8] optim: mean_diagonal applies the operator where it was built; CG(device='cuda') sees cuda:0 as its card A host-built operator given a card probe segfaulted in BART's normal, and a solve placed on 'cuda' recursed since cuda:0 != cuda. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01MFB6QNUxWSoR8RjtuWacSY --- src/bartorch/optim/_linear.py | 26 +++++++++++++++++++++----- tests/test_solve.py | 9 +++++++++ 2 files changed, 30 insertions(+), 5 deletions(-) diff --git a/src/bartorch/optim/_linear.py b/src/bartorch/optim/_linear.py index 6986f2bda..975f8d0dd 100644 --- a/src/bartorch/optim/_linear.py +++ b/src/bartorch/optim/_linear.py @@ -255,9 +255,10 @@ def mean_diagonal(A, *, probes: int = 1, seed: int = 0, device=None) -> float: Seed of the generator the probes are drawn from. device : torch.device or str, default=None Where the probes are drawn, applied and reduced; on the host when - omitted. On the operator's card a volume's probe never crosses the - bus, and the generator is the card's, so the estimate differs from the - host's by the probe's spread. + omitted. On the card the operator was built on, a volume's probe never + crosses the bus; the generator is the card's, so the estimate differs + from the host's by the probe's spread. An operator built elsewhere is + applied where it was built. Returns ------- @@ -271,14 +272,29 @@ def mean_diagonal(A, *, probes: int = 1, seed: int = 0, device=None) -> float: if probes < 1: raise ValueError(f"the estimate takes at least one probe, not {probes}") on = torch.device("cpu" if device is None else device) + home = torch.device(getattr(A, "device", None) or "cpu") generator = torch.Generator(device=on).manual_seed(int(seed)) total = 0.0 for _ in range(int(probes)): v = torch.randn(A.ishape, dtype=torch.complex64, generator=generator, device=on) - total += float(torch.vdot(v.reshape(-1), A.normal(v).reshape(-1)).real / v.norm() ** 2) + # A BART operator applies where it was built; a probe drawn elsewhere visits it. + made = A.normal(v.to(home)).to(on) + total += float(torch.vdot(v.reshape(-1), made.reshape(-1)).real / v.norm() ** 2) return total / probes +def _same_device(a: torch.device, b: torch.device) -> bool: + """Whether ``a`` and ``b`` are one device; ``cuda`` without an index is the current card.""" + if a.type != b.type: + return False + if a.type == "cuda": + current = torch.cuda.current_device() + return (a.index if a.index is not None else current) == ( + b.index if b.index is not None else current + ) + return a.index == b.index or a.index is None or b.index is None + + def _conditioner(precond, shape): """``precond`` as a BART operator, checked to map the image to itself.""" if precond is None: @@ -886,7 +902,7 @@ def _from_rhs(self, normal, b: torch.Tensor, x0: torch.Tensor | None) -> torch.T ) def _moves(self, y: torch.Tensor) -> bool: - return self.device is not None and y.device != self.device + return self.device is not None and not _same_device(y.device, self.device) def _placed(self, solve, y: torch.Tensor, x0: torch.Tensor | None) -> torch.Tensor: """``solve(y, x0)`` with both on :attr:`device`, the image brought back to ``y``'s.""" diff --git a/tests/test_solve.py b/tests/test_solve.py index e7e029710..3244ce92a 100644 --- a/tests/test_solve.py +++ b/tests/test_solve.py @@ -1040,6 +1040,15 @@ def test_a_device_where_the_data_already_is_changes_nothing(_small_encoding): assert "device='cpu'" in repr(optim.CG(device="cpu")) +def test_a_device_without_an_index_is_the_same_device(): + from bartorch.optim._linear import _same_device + + cpu = torch.device("cpu") + assert _same_device(cpu, torch.device("cpu")) + assert not _same_device(cpu, torch.device("cuda")) + assert not optim.CG(0.1, device="cpu")._moves(torch.zeros(2)) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a card") def test_the_vectors_of_a_solve_on_a_card_come_back_to_the_host(_small_encoding): A, y = _small_encoding From 222678d5f2a54ea9273958db7c24ee80af56112d Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 1 Nov 2026 23:32:02 +0000 Subject: [PATCH 6/8] optim: a card solve refuses a host-built encoding instead of crashing; scaling sorts once on a card CG(device='cuda') handed BART card vectors with a host operator and segfaulted; it now says to build the encoding on the card, as the test does. kthvalue over 67 M magnitudes took about a second on an A40; one radix sort gives the same order statistics. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01MFB6QNUxWSoR8RjtuWacSY --- src/bartorch/optim/_linear.py | 15 +++++++++++---- src/bartorch/optim/_scaling.py | 17 ++++++++++++----- tests/test_solve.py | 21 ++++++++++++++++----- 3 files changed, 39 insertions(+), 14 deletions(-) diff --git a/src/bartorch/optim/_linear.py b/src/bartorch/optim/_linear.py index 975f8d0dd..ce388c960 100644 --- a/src/bartorch/optim/_linear.py +++ b/src/bartorch/optim/_linear.py @@ -777,7 +777,7 @@ def __call__( """ from bartorch.linop._base import _tracking - if self._moves(y): + if self._moves(y, A): return self._placed(lambda b, w: self(b, A, w, steps=steps), y, x0) made = _item_by_item(self, y, A, x0, steps=steps) @@ -872,7 +872,7 @@ def inverse(self, b: torch.Tensor, A, x0: torch.Tensor | None = None) -> torch.T backward pass of a call solves through the same route, so it costs what the forward pass does. No gradient is recorded. """ - if self._moves(b): + if self._moves(b, A): return self._placed(lambda r, w: self.inverse(r, A, w), b, x0) normal = A.gram() for term in self.terms: @@ -901,8 +901,15 @@ def _from_rhs(self, normal, b: torch.Tensor, x0: torch.Tensor | None) -> torch.T **self._settings(), ) - def _moves(self, y: torch.Tensor) -> bool: - return self.device is not None and not _same_device(y.device, self.device) + def _moves(self, y: torch.Tensor, A=None) -> bool: + if self.device is None or _same_device(y.device, self.device): + return False + built = getattr(A, "device", None) + if self.device.type == "cuda" and (built is None or torch.device(built).type != "cuda"): + raise ValueError( + f"a solve on {self.device} needs an encoding built there; this one is on the host" + ) + return True def _placed(self, solve, y: torch.Tensor, x0: torch.Tensor | None) -> torch.Tensor: """``solve(y, x0)`` with both on :attr:`device`, the image brought back to ``y``'s.""" diff --git a/src/bartorch/optim/_scaling.py b/src/bartorch/optim/_scaling.py index baac97989..64fba050f 100644 --- a/src/bartorch/optim/_scaling.py +++ b/src/bartorch/optim/_scaling.py @@ -90,21 +90,28 @@ def data_scaling( return _scaling_norm(magnitudes, compat, percentile) -def _scaling_norm(magnitudes: torch.Tensor, compat: bool, percentile: float | None) -> float: +def _scaling_norm( + magnitudes: torch.Tensor, compat: bool, percentile: float | None, sort: bool | None = None +) -> float: """BART's ``estimate_scaling_norm`` over ``magnitudes``, by selection where they lie. The same order statistics BART sorts the whole array for -- the median at ``n // 2``, the 90th percentile at ``trunc(0.9 n)``, the maximum, or the - asked-for percentile -- found with ``kthvalue``, on the card when the - image is there. + asked-for percentile -- found with ``kthvalue`` on the host and by one + sort on a card, where selection over tens of millions of values is slow. """ n = magnitudes.numel() if n == 0: return 0.0 m = magnitudes.to(torch.float32) + # On a card one radix sort beats repeated selection; on the host selection wins. + ordered = torch.sort(m).values if (m.is_cuda if sort is None else sort) else None def at(i: int) -> float: - return float(torch.kthvalue(m, min(n - 1, max(0, i)) + 1).values) + i = min(n - 1, max(0, i)) + if ordered is not None: + return float(ordered[i]) + return float(torch.kthvalue(m, i + 1).values) if percentile is not None and percentile > 0: if percentile > 1: @@ -114,5 +121,5 @@ def at(i: int) -> float: if compat: return median p90 = at(int(n * 0.9)) - top = float(m.max()) + top = float(m.max() if ordered is None else ordered[-1]) return p90 if (top - p90) < 2 * (p90 - median) else top diff --git a/tests/test_solve.py b/tests/test_solve.py index 3244ce92a..e01ff7643 100644 --- a/tests/test_solve.py +++ b/tests/test_solve.py @@ -807,6 +807,11 @@ def test_the_scaling_of_nothing_is_zero_and_a_percentile_is_a_fraction(): assert _scaling_norm(torch.empty(0), False, None) == 0.0 with pytest.raises(ValueError, match="fraction"): _scaling_norm(torch.ones(4), False, 1.5) + m = torch.rand(1001, generator=torch.Generator().manual_seed(0)).square() + for compat, percentile in ((False, None), (True, None), (False, 0.99)): + assert _scaling_norm(m, compat, percentile, sort=True) == _scaling_norm( + m, compat, percentile, sort=False + ) def test_a_wavelet_term_reused_answers_as_a_freshly_built_one(): @@ -1040,23 +1045,29 @@ def test_a_device_where_the_data_already_is_changes_nothing(_small_encoding): assert "device='cpu'" in repr(optim.CG(device="cpu")) -def test_a_device_without_an_index_is_the_same_device(): +def test_a_device_without_an_index_is_the_same_device(_small_encoding): from bartorch.optim._linear import _same_device cpu = torch.device("cpu") assert _same_device(cpu, torch.device("cpu")) assert not _same_device(cpu, torch.device("cuda")) assert not optim.CG(0.1, device="cpu")._moves(torch.zeros(2)) + host, y = _small_encoding + with pytest.raises(ValueError, match="built there"): + optim.CG(0.1, device="cuda")(y, host) @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a card") def test_the_vectors_of_a_solve_on_a_card_come_back_to_the_host(_small_encoding): - A, y = _small_encoding + host, y = _small_encoding + torch.manual_seed(0) + A = linop.CartesianSense(_rand(3, 8, 8), (8, 8), device="cuda") on_card = optim.CG(0.1, maxiter=30, device="cuda") - b = A.adjoint(y) + b = host.adjoint(y) + want = optim.CG(0.1, maxiter=30)(y, host) for made in (on_card(y, A), on_card.inverse(b, A, x0=torch.zeros_like(b))): assert made.device == y.device - torch.testing.assert_close(made, optim.CG(0.1, maxiter=30)(y, A), rtol=1e-3, atol=1e-4) + torch.testing.assert_close(made, want, rtol=1e-3, atol=1e-4) def test_the_mean_diagonal_is_the_trace_over_the_size(_small_encoding): @@ -1103,7 +1114,7 @@ def test_a_solve_on_another_device_is_the_same_solve(_small_encoding, monkeypatc solver = optim.CG(0.1, maxiter=9, device="cpu") asked = [] - def moves(data): + def moves(data, A=None): asked.append(data.shape) return len(asked) == 1 From 74f88a6a1f74866fbfede79f6b9e966969fe8750 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 2 Nov 2026 02:14:42 +0000 Subject: [PATCH 7/8] tests: cover a placed solve's move and an index-free card without a card Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01MFB6QNUxWSoR8RjtuWacSY --- tests/test_solve.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/test_solve.py b/tests/test_solve.py index e01ff7643..361bb5c0b 100644 --- a/tests/test_solve.py +++ b/tests/test_solve.py @@ -1045,13 +1045,17 @@ def test_a_device_where_the_data_already_is_changes_nothing(_small_encoding): assert "device='cpu'" in repr(optim.CG(device="cpu")) -def test_a_device_without_an_index_is_the_same_device(_small_encoding): +def test_a_device_without_an_index_is_the_same_device(_small_encoding, monkeypatch): from bartorch.optim._linear import _same_device cpu = torch.device("cpu") assert _same_device(cpu, torch.device("cpu")) assert not _same_device(cpu, torch.device("cuda")) assert not optim.CG(0.1, device="cpu")._moves(torch.zeros(2)) + assert optim.CG(0.1, device="cpu")._moves(torch.zeros(2, device="meta")) + monkeypatch.setattr(torch.cuda, "current_device", lambda: 0) + assert _same_device(torch.device("cuda"), torch.device("cuda", 0)) + assert not _same_device(torch.device("cuda"), torch.device("cuda", 1)) host, y = _small_encoding with pytest.raises(ValueError, match="built there"): optim.CG(0.1, device="cuda")(y, host) From fdf8a8a8eb613d57cc2c280b34cc4965cdf24bac Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 2 Nov 2026 06:53:47 +0000 Subject: [PATCH 8/8] learning: Patchwise repeats a single item's keywords for every patch of a batch A class-conditioned network (MRI2Qmap's teacher) given one label for one volume failed its shape check at batch > 1. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01MFB6QNUxWSoR8RjtuWacSY --- src/bartorch/learning/_patches.py | 2 +- tests/test_learning_blocks.py | 16 ++++++++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/src/bartorch/learning/_patches.py b/src/bartorch/learning/_patches.py index c01a6d4a1..e766f46c0 100644 --- a/src/bartorch/learning/_patches.py +++ b/src/bartorch/learning/_patches.py @@ -266,7 +266,7 @@ def _run(self, pieces, count: int, n: int, sigma, keywords) -> torch.Tensor: """One call of the network over ``count`` patches of each of ``n`` items.""" def repeat(value): - if isinstance(value, torch.Tensor) and value.ndim and n == value.shape[0] and n > 1: + if isinstance(value, torch.Tensor) and value.ndim and n == value.shape[0]: return value.to(self.device).repeat(count) return value.to(self.device) if isinstance(value, torch.Tensor) else value diff --git a/tests/test_learning_blocks.py b/tests/test_learning_blocks.py index ff38a7752..782bd5333 100644 --- a/tests/test_learning_blocks.py +++ b/tests/test_learning_blocks.py @@ -473,6 +473,22 @@ def test_a_patchwise_network_is_called_without_a_noise_level_when_given_none(): assert "overlap=True" in repr(learning.Patchwise(net, (4, 4), device="cpu")) +class _Labelled(nn.Module): + """A network that, like a class-conditioned UNet, needs one label per patch.""" + + def forward(self, v, sigma, y): + assert y.shape == (v.shape[0],) and sigma.shape == (v.shape[0],) + return v + + +@pytest.mark.parametrize("batch", [1, 3]) +def test_one_items_keywords_are_repeated_for_every_patch_of_a_batch(batch): + run = learning.Patchwise(_Labelled(), (4, 4), device="cpu", shift=False, batch=batch) + x = torch.randn(1, 1, 8, 8) + made = run(x, torch.full((1,), 0.1), y=torch.zeros(1, dtype=torch.long)) + torch.testing.assert_close(made, x) + + def test_dropout_randomises_training_and_not_evaluation(): torch.manual_seed(0) net = _perturbed(learning.UNet(2, spatial=2, widths=(8, 16), dropout=0.5))