From 4424cb48890010b91ca1262c87e0852a34a7aacb Mon Sep 17 00:00:00 2001
From: Robert Smith <robert.smith@florey.edu.au>
Date: Mon, 4 Nov 2024 15:50:43 +1100
Subject: [PATCH 01/34] dwidenoise: Modularise kernel

Better separation of code responsible for fetching a batch of input data within a sliding spatial window from the code responsible for the denoising of the image data.
---
 cmd/dwidenoise.cpp | 218 ++++++++++++++++++++++++++++-----------------
 1 file changed, 136 insertions(+), 82 deletions(-)

diff --git a/cmd/dwidenoise.cpp b/cmd/dwidenoise.cpp
index 20d9242cd4..100c362007 100644
--- a/cmd/dwidenoise.cpp
+++ b/cmd/dwidenoise.cpp
@@ -144,29 +144,98 @@ void usage() {
 
 using real_type = float;
 
-template <typename F = float> class DenoisingFunctor {
+// Class to encode return information from kernel
+template <class MatrixType> class KernelData {
+public:
+  KernelData(const size_t volumes, const size_t kernel_size)
+      : centre_index(-1),                            //
+        voxel_count(kernel_size),                    //
+        X(MatrixType::Zero(volumes, kernel_size)) {} //
+  size_t centre_index;
+  size_t voxel_count;
+  MatrixType X;
+};
+
+template <class MatrixType> class KernelBase {
+public:
+  KernelBase() : pos({-1, -1, -1}) {}
+  KernelBase(const KernelBase &) : pos({-1, -1, -1}) {}
+
+protected:
+  // Store / restore position of image before / after data loading
+  std::array<ssize_t, 3> pos;
+  template <class ImageType> void stash_pos(const ImageType &image) {
+    for (size_t axis = 0; axis != 3; ++axis)
+      pos[axis] = image.index(axis);
+  }
+  template <class ImageType> void restore_pos(ImageType &image) {
+    for (size_t axis = 0; axis != 3; ++axis)
+      image.index(axis) = pos[axis];
+  }
+};
+
+template <class MatrixType> class KernelCube : public KernelBase<MatrixType> {
+public:
+  KernelCube(const std::vector<uint32_t> &extent)
+      : half_extent({int(extent[0] / 2), int(extent[1] / 2), int(extent[2] / 2)}) {
+    for (auto e : extent) {
+      if (!(e % 2))
+        throw Exception("Size of cubic kernel must be an odd integer");
+    }
+  }
+  KernelCube(const KernelCube &) = default;
+  template <class ImageType> void operator()(ImageType &image, KernelData<MatrixType> &data) {
+    assert(data.X.cols() == size());
+    KernelBase<MatrixType>::stash_pos(image);
+    size_t k = 0;
+    for (int z = -half_extent[2]; z <= half_extent[2]; z++) {
+      image.index(2) = wrapindex(z, 2, image.size(2));
+      for (int y = -half_extent[1]; y <= half_extent[1]; y++) {
+        image.index(1) = wrapindex(y, 1, image.size(1));
+        for (int x = -half_extent[0]; x <= half_extent[0]; x++, k++) {
+          image.index(0) = wrapindex(x, 0, image.size(0));
+          data.X.col(k) = image.row(3);
+        }
+      }
+    }
+    KernelBase<MatrixType>::restore_pos(image);
+    data.voxel_count = size();
+    data.centre_index = size() / 2;
+  }
+  size_t size() const { return (2 * half_extent[0] + 1) * (2 * half_extent[1] + 1) * (2 * half_extent[2] + 1); }
+
+private:
+  const std::vector<int> half_extent;
+
+  // patch handling at image edges
+  inline size_t wrapindex(int r, int axis, int max) const {
+    int rr = KernelBase<MatrixType>::pos[axis] + r;
+    if (rr < 0)
+      rr = half_extent[axis] - r;
+    if (rr >= max)
+      rr = (max - 1) - half_extent[axis] - r;
+    return rr;
+  }
+};
+
+template <typename F, class KernelType> class DenoisingFunctor {
 
 public:
   using MatrixType = Eigen::Matrix<F, Eigen::Dynamic, Eigen::Dynamic>;
   using SValsType = Eigen::VectorXd;
 
-  DenoisingFunctor(int ndwi,
-                   const std::vector<uint32_t> &extent,
-                   Image<bool> &mask,
-                   Image<real_type> &noise,
-                   Image<uint16_t> &rank,
-                   bool exp1)
-      : extent{{extent[0] / 2, extent[1] / 2, extent[2] / 2}},
+  DenoisingFunctor(
+      int ndwi, KernelType &kernel, Image<bool> &mask, Image<real_type> &noise, Image<uint16_t> &rank, bool exp1)
+      : data(ndwi, kernel.size()),
+        kernel(kernel),
         m(ndwi),
-        n(extent[0] * extent[1] * extent[2]),
+        n(kernel.size()),
         r(std::min(m, n)),
         q(std::max(m, n)),
         exp1(exp1),
-        X(m, n),
         XtX(r, r),
         eig(r),
         s(r),
-        pos{{0, 0, 0}},
         mask(mask),
         noise(noise),
         rankmap(rank) {}
@@ -180,13 +249,13 @@ template <typename F = float> class DenoisingFunctor {
     }
 
     // Load data in local window
-    load_data(dwi);
+    kernel(dwi, data);
 
     // Compute Eigendecomposition:
     if (m <= n)
-      XtX.template triangularView<Eigen::Lower>() = X * X.adjoint();
+      XtX.template triangularView<Eigen::Lower>() = data.X * data.X.adjoint();
     else
-      XtX.template triangularView<Eigen::Lower>() = X.adjoint() * X;
+      XtX.template triangularView<Eigen::Lower>() = data.X.adjoint() * data.X;
     eig.compute(XtX);
     // eigenvalues sorted in increasing order:
     s = eig.eigenvalues().template cast<double>();
@@ -215,14 +284,18 @@ template <typename F = float> class DenoisingFunctor {
       s.head(cutoff_p).setZero();
       s.tail(r - cutoff_p).setOnes();
       if (m <= n)
-        X.col(n / 2) = eig.eigenvectors() * (s.cast<F>().asDiagonal() * (eig.eigenvectors().adjoint() * X.col(n / 2)));
+        data.X.col(data.centre_index) =
+            eig.eigenvectors() *
+            (s.cast<F>().asDiagonal() * (eig.eigenvectors().adjoint() * data.X.col(data.centre_index)));
       else
-        X.col(n / 2) = X * (eig.eigenvectors() * (s.cast<F>().asDiagonal() * eig.eigenvectors().adjoint().col(n / 2)));
+        data.X.col(data.centre_index) =
+            data.X *
+            (eig.eigenvectors() * (s.cast<F>().asDiagonal() * eig.eigenvectors().adjoint().col(data.centre_index)));
     }
 
     // Store output
     assign_pos_of(dwi).to(out);
-    out.row(3) = X.col(n / 2);
+    out.row(3) = data.X.col(data.centre_index);
 
     // store noise map if requested:
     if (noise.valid()) {
@@ -237,60 +310,26 @@ template <typename F = float> class DenoisingFunctor {
   }
 
 private:
-  const std::array<ssize_t, 3> extent;
+  KernelData<MatrixType> data;
+  KernelType kernel;
   const ssize_t m, n, r, q;
   const bool exp1;
-  MatrixType X;
   MatrixType XtX;
   Eigen::SelfAdjointEigenSolver<MatrixType> eig;
   SValsType s;
-  std::array<ssize_t, 3> pos;
   double sigma2;
   Image<bool> mask;
   Image<real_type> noise;
   Image<uint16_t> rankmap;
-
-  template <typename ImageType> void load_data(ImageType &dwi) {
-    pos[0] = dwi.index(0);
-    pos[1] = dwi.index(1);
-    pos[2] = dwi.index(2);
-    // fill patch
-    X.setZero();
-    size_t k = 0;
-    for (int z = -extent[2]; z <= extent[2]; z++) {
-      dwi.index(2) = wrapindex(z, 2, dwi.size(2));
-      for (int y = -extent[1]; y <= extent[1]; y++) {
-        dwi.index(1) = wrapindex(y, 1, dwi.size(1));
-        for (int x = -extent[0]; x <= extent[0]; x++, k++) {
-          dwi.index(0) = wrapindex(x, 0, dwi.size(0));
-          X.col(k) = dwi.row(3);
-        }
-      }
-    }
-    // reset image position
-    dwi.index(0) = pos[0];
-    dwi.index(1) = pos[1];
-    dwi.index(2) = pos[2];
-  }
-
-  inline size_t wrapindex(int r, int axis, int max) const {
-    // patch handling at image edges
-    int rr = pos[axis] + r;
-    if (rr < 0)
-      rr = extent[axis] - r;
-    if (rr >= max)
-      rr = (max - 1) - extent[axis] - r;
-    return rr;
-  }
 };
 
-template <typename T>
+template <typename T, class KernelType>
 void process_image(Header &data,
                    Image<bool> &mask,
                    Image<real_type> &noise,
                    Image<uint16_t> &rank,
                    const std::string &output_name,
-                   const std::vector<uint32_t> &extent,
+                   KernelType &kernel,
                    bool exp1) {
   auto input = data.get_image<T>().with_direct_io(3);
   // create output
@@ -298,24 +337,20 @@ void process_image(Header &data,
   header.datatype() = DataType::from<T>();
   auto output = Image<T>::create(output_name, header);
   // run
-  DenoisingFunctor<T> func(data.size(3), extent, mask, noise, rank, exp1);
+  DenoisingFunctor<T, KernelType> func(data.size(3), kernel, mask, noise, rank, exp1);
   ThreadedLoop("running MP-PCA denoising", data, 0, 3).run(func, input, output);
 }
 
-void run() {
-  auto dwi = Header::open(argument[0]);
-
-  if (dwi.ndim() != 4 || dwi.size(3) <= 1)
-    throw Exception("input image must be 4-dimensional");
-
-  Image<bool> mask;
-  auto opt = get_options("mask");
-  if (!opt.empty()) {
-    mask = Image<bool>::open(opt[0][0]);
-    check_dimensions(mask, dwi, 0, 3);
-  }
-
-  opt = get_options("extent");
+template <typename T>
+void make_kernel(Header &data,
+                 Image<bool> &mask,
+                 Image<real_type> &noise,
+                 Image<uint16_t> &rank,
+                 const std::string &output_name,
+                 bool exp1) {
+  using KernelType = KernelCube<Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>>;
+
+  auto opt = get_options("extent");
   std::vector<uint32_t> extent;
   if (!opt.empty()) {
     extent = parse_ints<uint32_t>(opt[0][0]);
@@ -326,25 +361,44 @@ void run() {
     for (int i = 0; i < 3; i++) {
       if (!(extent[i] & 1))
         throw Exception("-extent must be a (list of) odd numbers");
-      if (extent[i] > dwi.size(i))
+      if (extent[i] > data.size(i))
         throw Exception("-extent must not exceed the image dimensions");
     }
   } else {
     uint32_t e = 1;
-    while (e * e * e < dwi.size(3))
+    while (Math::pow3(e) < data.size(3))
       e += 2;
-    extent = {
-        std::min(e, uint32_t(dwi.size(0))), std::min(e, uint32_t(dwi.size(1))), std::min(e, uint32_t(dwi.size(2)))};
+    extent = {std::min(e, uint32_t(data.size(0))),  //
+              std::min(e, uint32_t(data.size(1))),  //
+              std::min(e, uint32_t(data.size(2)))}; //
   }
   INFO("selected patch size: " + str(extent[0]) + " x " + str(extent[1]) + " x " + str(extent[2]) + ".");
 
-  bool exp1 = get_option_value("estimator", 1) == 0; // default: Exp2 (unbiased estimator)
+  if (std::min<uint32_t>(data.size(3), extent[0] * extent[1] * extent[2]) < 15) {
+    WARN("The number of volumes or the patch size is small. "
+         "This may lead to discretisation effects in the noise level "
+         "and cause inconsistent denoising between adjacent voxels.");
+  }
+
+  KernelType kernel(extent);
+  process_image<T, KernelType>(data, mask, noise, rank, output_name, kernel, exp1);
+}
 
-  if (std::min<uint32_t>(dwi.size(3), extent[0] * extent[1] * extent[2]) < 15) {
-    WARN("The number of volumes or the patch size is small. This may lead to discretisation effects "
-         "in the noise level and cause inconsistent denoising between adjacent voxels.");
+void run() {
+  auto dwi = Header::open(argument[0]);
+
+  if (dwi.ndim() != 4 || dwi.size(3) <= 1)
+    throw Exception("input image must be 4-dimensional");
+
+  Image<bool> mask;
+  auto opt = get_options("mask");
+  if (!opt.empty()) {
+    mask = Image<bool>::open(opt[0][0]);
+    check_dimensions(mask, dwi, 0, 3);
   }
 
+  bool exp1 = get_option_value("estimator", 1) == 0; // default: Exp2 (unbiased estimator)
+
   Image<real_type> noise;
   opt = get_options("noise");
   if (!opt.empty()) {
@@ -370,19 +424,19 @@ void run() {
   switch (prec) {
   case 0:
     INFO("select real float32 for processing");
-    process_image<float>(dwi, mask, noise, rank, argument[1], extent, exp1);
+    make_kernel<float>(dwi, mask, noise, rank, argument[1], exp1);
     break;
   case 1:
     INFO("select real float64 for processing");
-    process_image<double>(dwi, mask, noise, rank, argument[1], extent, exp1);
+    make_kernel<double>(dwi, mask, noise, rank, argument[1], exp1);
     break;
   case 2:
     INFO("select complex float32 for processing");
-    process_image<cfloat>(dwi, mask, noise, rank, argument[1], extent, exp1);
+    make_kernel<cfloat>(dwi, mask, noise, rank, argument[1], exp1);
     break;
   case 3:
     INFO("select complex float64 for processing");
-    process_image<cdouble>(dwi, mask, noise, rank, argument[1], extent, exp1);
+    make_kernel<cdouble>(dwi, mask, noise, rank, argument[1], exp1);
     break;
   }
 }

From bf0f978f348cf09d1018700d3ebed0d729a3dab3 Mon Sep 17 00:00:00 2001
From: Robert Smith <robert.smith@florey.edu.au>
Date: Tue, 5 Nov 2024 12:53:07 +1100
Subject: [PATCH 02/34] dwidenoise: First working version of spherical kernel

New default behaviour is to use an expanding spherical kernel with number of voxels at least 1.1 times the number of volumes. For voxels near the edge of the image FoV, the radius of the kernel will increase until the requisite number of voxels is obtained.
Note that execution speed of this implementation seems to be reduced, even when using the cuboid kernel; this may be due to use of Eigen Blocks to denoise voxels with kernels smaller than the maximum processed.
---
 cmd/dwidenoise.cpp | 325 ++++++++++++++++++++++++++++++++++-----------
 1 file changed, 248 insertions(+), 77 deletions(-)

diff --git a/cmd/dwidenoise.cpp b/cmd/dwidenoise.cpp
index 100c362007..e6538e18ea 100644
--- a/cmd/dwidenoise.cpp
+++ b/cmd/dwidenoise.cpp
@@ -26,6 +26,10 @@ using namespace App;
 const std::vector<std::string> dtypes = {"float32", "float64"};
 const std::vector<std::string> estimators = {"exp1", "exp2"};
 
+const std::vector<std::string> shapes = {"cuboid", "sphere"};
+enum class shape_type { CUBOID, SPHERE };
+constexpr default_type sphere_multiplier_default = 1.1;
+
 // clang-format off
 void usage() {
 
@@ -49,7 +53,29 @@ void usage() {
   + "Note that this function does not correct for non-Gaussian noise biases"
     " present in magnitude-reconstructed MRI images."
     " If available, including the MRI phase data can reduce such non-Gaussian biases,"
-    " and the command now supports complex input data.";
+    " and the command now supports complex input data."
+
+  + "The sliding spatial window behaves differently at the edges of the image FoV "
+    "depending on the shape / size selected for that window. "
+    "The default behaviour is to use a spherical kernel centred at the voxel of interest, "
+    "whose size is some multiple of the number of input volumes; "
+    "where some such voxels lie outside of the image FoV, "
+    "the radius of the kernel will be increased until the requisite number of voxels are used. "
+    "For a spherical kernel of a fixed radius, "
+    "no such expansion will occur, "
+    "and so for voxels near the image edge a reduced number of voxels will be present in the kernel. "
+    "For a cuboid kernel, "
+    "the centre of the kernel will be offset from the voxel being processed "
+    "such that the entire volume of the kernel resides within the image FoV."
+
+  + "The size of the default spherical kernel is set to select a number of voxels that is "
+    "1.1 times the number of volumes in the input series. "
+    "If a cuboid kernel is requested, "
+    "but the -extent option is not specified, "
+    "the command will select the smallest isotropic patch size "
+    "that exceeds the number of DW images in the input data; "
+    "e.g., 5x5x5 for data with <= 125 DWI volumes, "
+    "7x7x7 for data with <= 343 DWI volumes, etc.";
 
   AUTHOR = "Daan Christiaens (daan.christiaens@kcl.ac.uk)"
            " and Jelle Veraart (jelle.veraart@nyumc.org)"
@@ -73,19 +99,25 @@ void usage() {
   + Argument("out", "the output denoised DWI image.").type_image_out();
 
   OPTIONS
+  + OptionGroup("Options for modifying the application of PCA denoising")
   + Option("mask",
            "Only process voxels within the specified binary brain mask image.")
     + Argument("image").type_image_in()
+  + Option("datatype",
+           "Datatype for the eigenvalue decomposition"
+           " (single or double precision). "
+           "For complex input data,"
+           " this will select complex float32 or complex float64 datatypes.")
+    + Argument("float32/float64").type_choice(dtypes)
+  + Option("estimator",
+           "Select the noise level estimator"
+           " (default = Exp2),"
+           " either: \n"
+           "* Exp1: the original estimator used in Veraart et al. (2016), or \n"
+           "* Exp2: the improved estimator introduced in Cordero-Grande et al. (2019).")
+    + Argument("Exp1/Exp2").type_choice(estimators)
 
-  + Option("extent",
-           "Set the patch size of the denoising filter."
-           " By default, the command will select the smallest isotropic patch size"
-           " that exceeds the number of DW images in the input data,"
-           " e.g., 5x5x5 for data with <= 125 DWI volumes,"
-           " 7x7x7 for data with <= 343 DWI volumes,"
-           " etc.")
-    + Argument("window").type_sequence_int()
-
+  + OptionGroup("Options for exporting additional data regarding PCA behaviour")
   + Option("noise",
            "The output noise map,"
            " i.e., the estimated noise level 'sigma' in the data."
@@ -93,25 +125,33 @@ void usage() {
            " this will be the total noise level across real and imaginary channels,"
            " so a scale factor sqrt(2) applies.")
     + Argument("level").type_image_out()
-
   + Option("rank",
            "The selected signal rank of the output denoised image.")
     + Argument("cutoff").type_image_out()
 
-  + Option("datatype",
-           "Datatype for the eigenvalue decomposition"
-           " (single or double precision). "
-           "For complex input data,"
-           " this will select complex float32 or complex float64 datatypes.")
-    + Argument("float32/float64").type_choice(dtypes)
+  + OptionGroup("Options for controlling the sliding spatial window")
+  + Option("shape",
+           "Set the shape of the sliding spatial window. "
+           "Options are: " + join(shapes, ",") + "; default: sphere")
+    + Argument("choice").type_choice(shapes)
+  + Option("radius_mm",
+           "Set an absolute spherical kernel radius in mm")
+    + Argument("value").type_float(0.0)
+  + Option("radius_ratio",
+           "Set the spherical kernel radius as a ratio of number of input volumes "
+           "(default: 1.1)")
+    + Argument("value").type_float(0.0)
+
+  + Option("extent",
+           "Set the patch size of the cuboid filter; "
+           "can be either a single odd integer or a comma-separated triplet of odd integers")
+    + Argument("window").type_sequence_int();
+
+
+
+
+
 
-  + Option("estimator",
-           "Select the noise level estimator"
-           " (default = Exp2),"
-           " either: \n"
-           "* Exp1: the original estimator used in Veraart et al. (2016), or \n"
-           "* Exp2: the improved estimator introduced in Cordero-Grande et al. (2019).")
-    + Argument("Exp1/Exp2").type_choice(estimators);
 
   COPYRIGHT =
       "Copyright (c) 2016 New York University, University of Antwerp, and the MRtrix3 contributors \n \n"
@@ -160,6 +200,7 @@ template <class MatrixType> class KernelBase {
 public:
   KernelBase() : pos({-1, -1, -1}) {}
   KernelBase(const KernelBase &) : pos({-1, -1, -1}) {}
+  virtual ssize_t minimum_size() const = 0;
 
 protected:
   // Store / restore position of image before / after data loading
@@ -185,7 +226,7 @@ template <class MatrixType> class KernelCube : public KernelBase<MatrixType> {
   }
   KernelCube(const KernelCube &) = default;
   template <class ImageType> void operator()(ImageType &image, KernelData<MatrixType> &data) {
-    assert(data.X.cols() == size());
+    assert(data.X.cols() == minimum_size());
     KernelBase<MatrixType>::stash_pos(image);
     size_t k = 0;
     for (int z = -half_extent[2]; z <= half_extent[2]; z++) {
@@ -199,10 +240,12 @@ template <class MatrixType> class KernelCube : public KernelBase<MatrixType> {
       }
     }
     KernelBase<MatrixType>::restore_pos(image);
-    data.voxel_count = size();
-    data.centre_index = size() / 2;
+    data.voxel_count = minimum_size();
+    data.centre_index = minimum_size() / 2;
+  }
+  ssize_t minimum_size() const override {
+    return (2 * half_extent[0] + 1) * (2 * half_extent[1] + 1) * (2 * half_extent[2] + 1);
   }
-  size_t size() const { return (2 * half_extent[0] + 1) * (2 * half_extent[1] + 1) * (2 * half_extent[2] + 1); }
 
 private:
   const std::vector<int> half_extent;
@@ -218,6 +261,106 @@ template <class MatrixType> class KernelCube : public KernelBase<MatrixType> {
   }
 };
 
+template <class MatrixType> class KernelSphereBase : public KernelBase<MatrixType> {
+public:
+  KernelSphereBase(const Header &voxel_grid, const default_type max_radius)
+      : shared(new Shared(voxel_grid, max_radius)) {}
+
+protected:
+  class Shared {
+  public:
+    using MapType = std::multimap<float, std::array<int, 3>>;
+    Shared(const Header &voxel_grid, const default_type max_radius) {
+      const default_type max_radius_sq = Math::pow2(max_radius);
+      const std::array<int, 3> half_extents({int(std::ceil(max_radius / voxel_grid.spacing(0))),   //
+                                             int(std::ceil(max_radius / voxel_grid.spacing(1))),   //
+                                             int(std::ceil(max_radius / voxel_grid.spacing(2)))}); //
+      // Build the searchlight
+      std::array<int, 3> offset;
+      for (offset[2] = -half_extents[2]; offset[2] <= half_extents[2]; ++offset[2]) {
+        for (offset[1] = -half_extents[1]; offset[1] <= half_extents[1]; ++offset[1]) {
+          for (offset[0] = -half_extents[0]; offset[0] <= half_extents[0]; ++offset[0]) {
+            const default_type squared_distance = Math::pow2(offset[0] * voxel_grid.spacing(0))    //
+                                                  + Math::pow2(offset[1] * voxel_grid.spacing(1))  //
+                                                  + Math::pow2(offset[2] * voxel_grid.spacing(2)); //
+            if (squared_distance <= max_radius_sq)
+              data.insert({squared_distance, offset});
+          }
+        }
+      }
+    }
+    MapType::const_iterator begin() const { return data.begin(); }
+    MapType::const_iterator end() const { return data.end(); }
+
+  private:
+    MapType data;
+  };
+  std::shared_ptr<Shared> shared;
+};
+
+template <class MatrixType> class KernelSphereRatio : public KernelSphereBase<MatrixType> {
+public:
+  KernelSphereRatio(const Header &voxel_grid, const default_type min_ratio)
+      : KernelSphereBase<MatrixType>(voxel_grid, compute_max_radius(voxel_grid, min_ratio)),
+        min_size(std::ceil(voxel_grid.size(3) * min_ratio)) {}
+  template <class ImageType> void operator()(ImageType &image, KernelData<MatrixType> &data) {
+    KernelBase<MatrixType>::stash_pos(image);
+    data.voxel_count = 0;
+    default_type prev_distance = -std::numeric_limits<default_type>::infinity();
+    auto map_it = KernelSphereBase<MatrixType>::shared->begin();
+    while (map_it != KernelSphereBase<MatrixType>::shared->end()) {
+      // If there's a tie in distances, want to include all such offsets in the kernel,
+      //   even if the size of the utilised kernel extends beyond the minimum size
+      if (map_it->first != prev_distance && data.voxel_count >= min_size)
+        break;
+      for (size_t axis = 0; axis != 3; ++axis)
+        image.index(axis) = KernelBase<MatrixType>::pos[axis] + map_it->second[axis];
+      if (!is_out_of_bounds(image, 0, 3)) {
+        // Is this larger than any kernel this thread has previously encountered?
+        // If so, try to project what the final size is going to be,
+        //   based on the set of voxels with identical distance to this one
+        //   all getting included in the kernel
+        if (data.voxel_count == data.X.cols()) {
+          size_t extra_cols = 1;
+          auto forward_search = map_it;
+          for (++forward_search;
+               forward_search != KernelSphereBase<MatrixType>::shared->end() && forward_search->first == map_it->first;
+               ++forward_search)
+            ++extra_cols;
+          data.X.conservativeResize(data.X.rows(), data.voxel_count + extra_cols);
+        }
+        data.X.col(data.voxel_count) = image.row(3);
+        prev_distance = map_it->first;
+        ++data.voxel_count;
+      }
+      ++map_it;
+    }
+    if (map_it == KernelSphereBase<MatrixType>::shared->end())
+      throw Exception("Inadequate spherical kernel initialisation");
+    KernelBase<MatrixType>::restore_pos(image);
+    data.centre_index = 0;
+  }
+  ssize_t minimum_size() const override { return min_size; }
+
+private:
+  size_t min_size;
+  // Determine an appropriate bounding box from which to generate the search table
+  // Find the radius for which 7/8 of the sphere will contain the minimum number of voxels, then round up
+  // This is only for setting the maximal radius for generation of the lookup table
+  default_type compute_max_radius(const Header &voxel_grid, const default_type min_ratio) const {
+    const size_t num_volumes = voxel_grid.size(3);
+    const default_type voxel_volume = voxel_grid.spacing(0) * voxel_grid.spacing(1) * voxel_grid.spacing(2);
+    const default_type sphere_volume = 8.0 * num_volumes * min_ratio * voxel_volume;
+    const default_type approx_radius = std::sqrt(sphere_volume * 0.75 / Math::pi);
+    const std::array<int, 3> half_extents({int(std::ceil(approx_radius / voxel_grid.spacing(0))),   //
+                                           int(std::ceil(approx_radius / voxel_grid.spacing(1))),   //
+                                           int(std::ceil(approx_radius / voxel_grid.spacing(2)))}); //
+    return std::max({half_extents[0] * voxel_grid.spacing(0),
+                     half_extents[1] * voxel_grid.spacing(1),
+                     half_extents[2] * voxel_grid.spacing(2)});
+  }
+};
+
 template <typename F, class KernelType> class DenoisingFunctor {
 
 public:
@@ -226,16 +369,13 @@ template <typename F, class KernelType> class DenoisingFunctor {
 
   DenoisingFunctor(
       int ndwi, KernelType &kernel, Image<bool> &mask, Image<real_type> &noise, Image<uint16_t> &rank, bool exp1)
-      : data(ndwi, kernel.size()),
+      : data(ndwi, kernel.minimum_size()),
         kernel(kernel),
         m(ndwi),
-        n(kernel.size()),
-        r(std::min(m, n)),
-        q(std::max(m, n)),
         exp1(exp1),
-        XtX(r, r),
-        eig(r),
-        s(r),
+        XtX(std::min(m, kernel.minimum_size()), std::min(m, kernel.minimum_size())),
+        eig(std::min(m, kernel.minimum_size())),
+        s(std::min(m, kernel.minimum_size())),
         mask(mask),
         noise(noise),
         rankmap(rank) {}
@@ -250,15 +390,28 @@ template <typename F, class KernelType> class DenoisingFunctor {
 
     // Load data in local window
     kernel(dwi, data);
+    auto X = data.X.leftCols(data.voxel_count);
+
+    const ssize_t n = data.voxel_count;
+    const ssize_t r = std::min(m, n);
+    const ssize_t q = std::max(m, n);
+
+    if (r > XtX.rows()) {
+      XtX.resize(r, r);
+      s.resize(r);
+    }
+
+    // TODO Fill matrices with NaN when in debug mode;
+    //   make sure results from one voxel are not creeping into another
 
     // Compute Eigendecomposition:
     if (m <= n)
-      XtX.template triangularView<Eigen::Lower>() = data.X * data.X.adjoint();
+      XtX.topLeftCorner(r, r).template triangularView<Eigen::Lower>() = X * X.adjoint();
     else
-      XtX.template triangularView<Eigen::Lower>() = data.X.adjoint() * data.X;
-    eig.compute(XtX);
+      XtX.topLeftCorner(r, r).template triangularView<Eigen::Lower>() = X.adjoint() * X;
+    eig.compute(XtX.topLeftCorner(r, r));
     // eigenvalues sorted in increasing order:
-    s = eig.eigenvalues().template cast<double>();
+    s.head(r) = eig.eigenvalues().template cast<double>();
 
     // Marchenko-Pastur optimal threshold
     const double lam_r = std::max(s[0], 0.0) / q;
@@ -282,20 +435,18 @@ template <typename F, class KernelType> class DenoisingFunctor {
     if (cutoff_p > 0) {
       // recombine data using only eigenvectors above threshold:
       s.head(cutoff_p).setZero();
-      s.tail(r - cutoff_p).setOnes();
+      s.segment(cutoff_p, r - cutoff_p).setOnes();
       if (m <= n)
-        data.X.col(data.centre_index) =
-            eig.eigenvectors() *
-            (s.cast<F>().asDiagonal() * (eig.eigenvectors().adjoint() * data.X.col(data.centre_index)));
+        X.col(data.centre_index) = eig.eigenvectors() * (s.head(r).cast<F>().asDiagonal() *
+                                                         (eig.eigenvectors().adjoint() * X.col(data.centre_index)));
       else
-        data.X.col(data.centre_index) =
-            data.X *
-            (eig.eigenvectors() * (s.cast<F>().asDiagonal() * eig.eigenvectors().adjoint().col(data.centre_index)));
+        X.col(data.centre_index) = X * (eig.eigenvectors() * (s.head(r).cast<F>().asDiagonal() *
+                                                              eig.eigenvectors().adjoint().col(data.centre_index)));
     }
 
     // Store output
     assign_pos_of(dwi).to(out);
-    out.row(3) = data.X.col(data.centre_index);
+    out.row(3) = X.col(data.centre_index);
 
     // store noise map if requested:
     if (noise.valid()) {
@@ -312,7 +463,7 @@ template <typename F, class KernelType> class DenoisingFunctor {
 private:
   KernelData<MatrixType> data;
   KernelType kernel;
-  const ssize_t m, n, r, q;
+  const ssize_t m;
   const bool exp1;
   MatrixType XtX;
   Eigen::SelfAdjointEigenSolver<MatrixType> eig;
@@ -348,40 +499,60 @@ void make_kernel(Header &data,
                  Image<uint16_t> &rank,
                  const std::string &output_name,
                  bool exp1) {
-  using KernelType = KernelCube<Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>>;
+  using MatrixType = Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>;
 
-  auto opt = get_options("extent");
-  std::vector<uint32_t> extent;
-  if (!opt.empty()) {
-    extent = parse_ints<uint32_t>(opt[0][0]);
-    if (extent.size() == 1)
-      extent = {extent[0], extent[0], extent[0]};
-    if (extent.size() != 3)
-      throw Exception("-extent must be either a scalar or a list of length 3");
-    for (int i = 0; i < 3; i++) {
-      if (!(extent[i] & 1))
-        throw Exception("-extent must be a (list of) odd numbers");
-      if (extent[i] > data.size(i))
-        throw Exception("-extent must not exceed the image dimensions");
+  auto opt = get_options("shape");
+  const shape_type shape = opt.empty() ? shape_type::SPHERE : shape_type((int)(opt[0][0]));
+
+  switch (shape) {
+  case shape_type::SPHERE: {
+    // TODO Could infer that user wants a cuboid kernel if -extent is used, even if -shape is not
+    if (!get_options("extent").empty()) {
+      throw Exception("-extent option does not apply to spherical kernel");
     }
-  } else {
-    uint32_t e = 1;
-    while (Math::pow3(e) < data.size(3))
-      e += 2;
-    extent = {std::min(e, uint32_t(data.size(0))),  //
-              std::min(e, uint32_t(data.size(1))),  //
-              std::min(e, uint32_t(data.size(2)))}; //
+    const default_type min_ratio = get_option_value("-radius_ratio", sphere_multiplier_default);
+    KernelSphereRatio<MatrixType> kernel(data, min_ratio);
+    process_image<T, KernelSphereRatio<MatrixType>>(data, mask, noise, rank, output_name, kernel, exp1);
+    return;
   }
-  INFO("selected patch size: " + str(extent[0]) + " x " + str(extent[1]) + " x " + str(extent[2]) + ".");
+  case shape_type::CUBOID: {
+    opt = get_options("extent");
+    std::vector<uint32_t> extent;
+    if (!opt.empty()) {
+      extent = parse_ints<uint32_t>(opt[0][0]);
+      if (extent.size() == 1)
+        extent = {extent[0], extent[0], extent[0]};
+      if (extent.size() != 3)
+        throw Exception("-extent must be either a scalar or a list of length 3");
+      for (int i = 0; i < 3; i++) {
+        if (!(extent[i] & 1))
+          throw Exception("-extent must be a (list of) odd numbers");
+        if (extent[i] > data.size(i))
+          throw Exception("-extent must not exceed the image dimensions");
+      }
+    } else {
+      uint32_t e = 1;
+      while (Math::pow3(e) < data.size(3))
+        e += 2;
+      extent = {std::min(e, uint32_t(data.size(0))),  //
+                std::min(e, uint32_t(data.size(1))),  //
+                std::min(e, uint32_t(data.size(2)))}; //
+    }
+    INFO("selected patch size: " + str(extent[0]) + " x " + str(extent[1]) + " x " + str(extent[2]) + ".");
 
-  if (std::min<uint32_t>(data.size(3), extent[0] * extent[1] * extent[2]) < 15) {
-    WARN("The number of volumes or the patch size is small. "
-         "This may lead to discretisation effects in the noise level "
-         "and cause inconsistent denoising between adjacent voxels.");
-  }
+    if (std::min<uint32_t>(data.size(3), extent[0] * extent[1] * extent[2]) < 15) {
+      WARN("The number of volumes or the patch size is small. "
+           "This may lead to discretisation effects in the noise level "
+           "and cause inconsistent denoising between adjacent voxels.");
+    }
 
-  KernelType kernel(extent);
-  process_image<T, KernelType>(data, mask, noise, rank, output_name, kernel, exp1);
+    KernelCube<MatrixType> kernel(extent);
+    process_image<T, KernelCube<MatrixType>>(data, mask, noise, rank, output_name, kernel, exp1);
+    return;
+  }
+  default:
+    assert(false);
+  }
 }
 
 void run() {

From 02c18b5a54c5bcbf5fea65ce4ee9131f312a1d37 Mon Sep 17 00:00:00 2001
From: Robert Smith <robert.smith@florey.edu.au>
Date: Tue, 5 Nov 2024 13:23:13 +1100
Subject: [PATCH 03/34] dwidenoise: Further changes to kernels

- Added ability to define a spherical kernel of a fixed radius. This will result in voxels near the edges of the image FoV having fewer voxels within the PCA kernel.
- Added option -voxels, to generate a spatial map of the number of voxels utilised in the PCA decomposition at each voxel location.
---
 cmd/dwidenoise.cpp | 146 ++++++++++++++++++++++++++++++++-------------
 1 file changed, 104 insertions(+), 42 deletions(-)

diff --git a/cmd/dwidenoise.cpp b/cmd/dwidenoise.cpp
index e6538e18ea..ebf3f16514 100644
--- a/cmd/dwidenoise.cpp
+++ b/cmd/dwidenoise.cpp
@@ -128,6 +128,9 @@ void usage() {
   + Option("rank",
            "The selected signal rank of the output denoised image.")
     + Argument("cutoff").type_image_out()
+  + Option("voxels",
+           "The number of voxels that contributed to the PCA for processing of each voxel")
+    + Argument("image").type_image_out()
 
   + OptionGroup("Options for controlling the sliding spatial window")
   + Option("shape",
@@ -141,18 +144,11 @@ void usage() {
            "Set the spherical kernel radius as a ratio of number of input volumes "
            "(default: 1.1)")
     + Argument("value").type_float(0.0)
-
   + Option("extent",
            "Set the patch size of the cuboid filter; "
            "can be either a single odd integer or a comma-separated triplet of odd integers")
     + Argument("window").type_sequence_int();
 
-
-
-
-
-
-
   COPYRIGHT =
       "Copyright (c) 2016 New York University, University of Antwerp, and the MRtrix3 contributors \n \n"
       "Permission is hereby granted, free of charge, to any non-commercial entity ('Recipient') obtaining a copy of "
@@ -200,7 +196,8 @@ template <class MatrixType> class KernelBase {
 public:
   KernelBase() : pos({-1, -1, -1}) {}
   KernelBase(const KernelBase &) : pos({-1, -1, -1}) {}
-  virtual ssize_t minimum_size() const = 0;
+  // This is just for pre-allocating matrices
+  virtual ssize_t estimated_size() const = 0;
 
 protected:
   // Store / restore position of image before / after data loading
@@ -226,7 +223,7 @@ template <class MatrixType> class KernelCube : public KernelBase<MatrixType> {
   }
   KernelCube(const KernelCube &) = default;
   template <class ImageType> void operator()(ImageType &image, KernelData<MatrixType> &data) {
-    assert(data.X.cols() == minimum_size());
+    assert(data.X.cols() == size());
     KernelBase<MatrixType>::stash_pos(image);
     size_t k = 0;
     for (int z = -half_extent[2]; z <= half_extent[2]; z++) {
@@ -240,10 +237,11 @@ template <class MatrixType> class KernelCube : public KernelBase<MatrixType> {
       }
     }
     KernelBase<MatrixType>::restore_pos(image);
-    data.voxel_count = minimum_size();
-    data.centre_index = minimum_size() / 2;
+    data.voxel_count = size();
+    data.centre_index = size() / 2;
   }
-  ssize_t minimum_size() const override {
+  ssize_t size() const { return estimated_size(); }
+  ssize_t estimated_size() const override {
     return (2 * half_extent[0] + 1) * (2 * half_extent[1] + 1) * (2 * half_extent[2] + 1);
   }
 
@@ -340,10 +338,10 @@ template <class MatrixType> class KernelSphereRatio : public KernelSphereBase<Ma
     KernelBase<MatrixType>::restore_pos(image);
     data.centre_index = 0;
   }
-  ssize_t minimum_size() const override { return min_size; }
+  ssize_t estimated_size() const override { return min_size; }
 
 private:
-  size_t min_size;
+  ssize_t min_size;
   // Determine an appropriate bounding box from which to generate the search table
   // Find the radius for which 7/8 of the sphere will contain the minimum number of voxels, then round up
   // This is only for setting the maximal radius for generation of the lookup table
@@ -361,24 +359,63 @@ template <class MatrixType> class KernelSphereRatio : public KernelSphereBase<Ma
   }
 };
 
+template <class MatrixType> class KernelSphereFixedRadius : public KernelSphereBase<MatrixType> {
+public:
+  KernelSphereFixedRadius(const Header &voxel_grid, const default_type radius)
+      : KernelSphereBase<MatrixType>(voxel_grid, radius),
+        maximum_size(std::distance(KernelSphereBase<MatrixType>::shared->begin(),  //
+                                   KernelSphereBase<MatrixType>::shared->end())) { //
+    INFO("Maximum number of voxels in " + str(radius) + "mm fixed-radius kernel is " + str(maximum_size));
+  }
+  template <class ImageType> void operator()(ImageType &image, KernelData<MatrixType> &data) {
+    KernelBase<MatrixType>::stash_pos(image);
+    data.voxel_count = 0;
+    default_type prev_distance = -std::numeric_limits<default_type>::infinity();
+    for (auto map_it = KernelSphereBase<MatrixType>::shared->begin();
+         map_it != KernelSphereBase<MatrixType>::shared->end();
+         ++map_it) {
+      for (size_t axis = 0; axis != 3; ++axis)
+        image.index(axis) = KernelBase<MatrixType>::pos[axis] + map_it->second[axis];
+      if (!is_out_of_bounds(image, 0, 3)) {
+        // We should not need to do any matrix size checking here;
+        //   it should have already been allocated to the maximum size of the kernel
+        data.X.col(data.voxel_count) = image.row(3);
+        ++data.voxel_count;
+      }
+    }
+    KernelBase<MatrixType>::restore_pos(image);
+    data.centre_index = 0;
+  }
+  ssize_t estimated_size() const override { return maximum_size; }
+
+private:
+  const ssize_t maximum_size;
+};
+
 template <typename F, class KernelType> class DenoisingFunctor {
 
 public:
   using MatrixType = Eigen::Matrix<F, Eigen::Dynamic, Eigen::Dynamic>;
   using SValsType = Eigen::VectorXd;
 
-  DenoisingFunctor(
-      int ndwi, KernelType &kernel, Image<bool> &mask, Image<real_type> &noise, Image<uint16_t> &rank, bool exp1)
-      : data(ndwi, kernel.minimum_size()),
+  DenoisingFunctor(int ndwi,
+                   KernelType &kernel,
+                   Image<bool> &mask,
+                   Image<real_type> &noise,
+                   Image<uint16_t> &rank,
+                   Image<uint16_t> &voxels,
+                   bool exp1)
+      : data(ndwi, kernel.estimated_size()),
         kernel(kernel),
         m(ndwi),
         exp1(exp1),
-        XtX(std::min(m, kernel.minimum_size()), std::min(m, kernel.minimum_size())),
-        eig(std::min(m, kernel.minimum_size())),
-        s(std::min(m, kernel.minimum_size())),
+        XtX(std::min(m, kernel.estimated_size()), std::min(m, kernel.estimated_size())),
+        eig(std::min(m, kernel.estimated_size())),
+        s(std::min(m, kernel.estimated_size())),
         mask(mask),
         noise(noise),
-        rankmap(rank) {}
+        rankmap(rank),
+        voxelsmap(voxels) {}
 
   template <typename ImageType> void operator()(ImageType &dwi, ImageType &out) {
     // Process voxels in mask only
@@ -458,6 +495,11 @@ template <typename F, class KernelType> class DenoisingFunctor {
       assign_pos_of(dwi, 0, 3).to(rankmap);
       rankmap.value() = uint16_t(r - cutoff_p);
     }
+    // store number of voxels map if requested:
+    if (voxelsmap.valid()) {
+      assign_pos_of(dwi, 0, 3).to(voxelsmap);
+      voxelsmap.value() = n;
+    }
   }
 
 private:
@@ -472,33 +514,36 @@ template <typename F, class KernelType> class DenoisingFunctor {
   Image<bool> mask;
   Image<real_type> noise;
   Image<uint16_t> rankmap;
+  Image<uint16_t> voxelsmap;
 };
 
 template <typename T, class KernelType>
-void process_image(Header &data,
-                   Image<bool> &mask,
-                   Image<real_type> &noise,
-                   Image<uint16_t> &rank,
-                   const std::string &output_name,
-                   KernelType &kernel,
-                   bool exp1) {
+void run(Header &data,
+         Image<bool> &mask,
+         Image<real_type> &noise,
+         Image<uint16_t> &rank,
+         Image<uint16_t> &voxels,
+         const std::string &output_name,
+         KernelType &kernel,
+         bool exp1) {
   auto input = data.get_image<T>().with_direct_io(3);
   // create output
   Header header(data);
   header.datatype() = DataType::from<T>();
   auto output = Image<T>::create(output_name, header);
   // run
-  DenoisingFunctor<T, KernelType> func(data.size(3), kernel, mask, noise, rank, exp1);
+  DenoisingFunctor<T, KernelType> func(data.size(3), kernel, mask, noise, rank, voxels, exp1);
   ThreadedLoop("running MP-PCA denoising", data, 0, 3).run(func, input, output);
 }
 
 template <typename T>
-void make_kernel(Header &data,
-                 Image<bool> &mask,
-                 Image<real_type> &noise,
-                 Image<uint16_t> &rank,
-                 const std::string &output_name,
-                 bool exp1) {
+void run(Header &data,
+         Image<bool> &mask,
+         Image<real_type> &noise,
+         Image<uint16_t> &rank,
+         Image<uint16_t> &voxels,
+         const std::string &output_name,
+         bool exp1) {
   using MatrixType = Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>;
 
   auto opt = get_options("shape");
@@ -507,15 +552,22 @@ void make_kernel(Header &data,
   switch (shape) {
   case shape_type::SPHERE: {
     // TODO Could infer that user wants a cuboid kernel if -extent is used, even if -shape is not
-    if (!get_options("extent").empty()) {
+    if (!get_options("extent").empty())
       throw Exception("-extent option does not apply to spherical kernel");
+    opt = get_options("radius_mm");
+    if (opt.size()) {
+      KernelSphereFixedRadius<MatrixType> kernel(data, opt[0][0]);
+      run<T, KernelSphereFixedRadius<MatrixType>>(data, mask, noise, rank, voxels, output_name, kernel, exp1);
+      return;
     }
     const default_type min_ratio = get_option_value("-radius_ratio", sphere_multiplier_default);
     KernelSphereRatio<MatrixType> kernel(data, min_ratio);
-    process_image<T, KernelSphereRatio<MatrixType>>(data, mask, noise, rank, output_name, kernel, exp1);
+    run<T, KernelSphereRatio<MatrixType>>(data, mask, noise, rank, voxels, output_name, kernel, exp1);
     return;
   }
   case shape_type::CUBOID: {
+    if (!get_options("radius_mm").size() || !get_options("radius_ratio").empty())
+      throw Exception("-radius_* options are inapplicable if cuboid kernel shape is selected");
     opt = get_options("extent");
     std::vector<uint32_t> extent;
     if (!opt.empty()) {
@@ -547,7 +599,7 @@ void make_kernel(Header &data,
     }
 
     KernelCube<MatrixType> kernel(extent);
-    process_image<T, KernelCube<MatrixType>>(data, mask, noise, rank, output_name, kernel, exp1);
+    run<T, KernelCube<MatrixType>>(data, mask, noise, rank, voxels, output_name, kernel, exp1);
     return;
   }
   default:
@@ -589,25 +641,35 @@ void run() {
     rank = Image<uint16_t>::create(opt[0][0], header);
   }
 
+  Image<uint16_t> voxels;
+  opt = get_options("voxels");
+  if (!opt.empty()) {
+    Header header(dwi);
+    header.ndim() = 3;
+    header.datatype() = DataType::UInt16;
+    header.reset_intensity_scaling();
+    voxels = Image<uint16_t>::create(opt[0][0], header);
+  }
+
   int prec = get_option_value("datatype", 0); // default: single precision
   if (dwi.datatype().is_complex())
     prec += 2; // support complex input data
   switch (prec) {
   case 0:
     INFO("select real float32 for processing");
-    make_kernel<float>(dwi, mask, noise, rank, argument[1], exp1);
+    run<float>(dwi, mask, noise, rank, voxels, argument[1], exp1);
     break;
   case 1:
     INFO("select real float64 for processing");
-    make_kernel<double>(dwi, mask, noise, rank, argument[1], exp1);
+    run<double>(dwi, mask, noise, rank, voxels, argument[1], exp1);
     break;
   case 2:
     INFO("select complex float32 for processing");
-    make_kernel<cfloat>(dwi, mask, noise, rank, argument[1], exp1);
+    run<cfloat>(dwi, mask, noise, rank, voxels, argument[1], exp1);
     break;
   case 3:
     INFO("select complex float64 for processing");
-    make_kernel<cdouble>(dwi, mask, noise, rank, argument[1], exp1);
+    run<cdouble>(dwi, mask, noise, rank, voxels, argument[1], exp1);
     break;
   }
 }

From 58d313d2d8c37c68f31d11e720affea6a84e4f1b Mon Sep 17 00:00:00 2001
From: Robert Smith <robert.smith@florey.edu.au>
Date: Tue, 5 Nov 2024 13:58:33 +1100
Subject: [PATCH 04/34] dwidenoise: Check validity of block operations

---
 cmd/dwidenoise.cpp | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/cmd/dwidenoise.cpp b/cmd/dwidenoise.cpp
index ebf3f16514..8cbf5c1bba 100644
--- a/cmd/dwidenoise.cpp
+++ b/cmd/dwidenoise.cpp
@@ -438,8 +438,14 @@ template <typename F, class KernelType> class DenoisingFunctor {
       s.resize(r);
     }
 
-    // TODO Fill matrices with NaN when in debug mode;
+    // Fill matrices with NaN when in debug mode;
     //   make sure results from one voxel are not creeping into another
+    //   due to use of block oberations to prevent memory re-allocation
+    //   in the presence of variation in kernel sizes
+#ifndef NDEBUG
+    XtX.fill(std::numeric_limits<F>::signaling_NaN());
+    s.fill(std::numeric_limits<default_type>::signaling_NaN());
+#endif
 
     // Compute Eigendecomposition:
     if (m <= n)

From 3c7645784715b82fdbc5d16361235f733791bc2c Mon Sep 17 00:00:00 2001
From: Robert Smith <robert.smith@florey.edu.au>
Date: Tue, 5 Nov 2024 15:45:08 +1100
Subject: [PATCH 05/34] dwidenoise: Un-template spatial kernels

---
 cmd/dwidenoise.cpp | 377 +++++++++++++++++++++------------------------
 1 file changed, 176 insertions(+), 201 deletions(-)

diff --git a/cmd/dwidenoise.cpp b/cmd/dwidenoise.cpp
index 8cbf5c1bba..7214150a9d 100644
--- a/cmd/dwidenoise.cpp
+++ b/cmd/dwidenoise.cpp
@@ -179,90 +179,82 @@ void usage() {
 // clang-format on
 
 using real_type = float;
+using voxel_type = Eigen::Array<int, 3, 1>;
 
 // Class to encode return information from kernel
-template <class MatrixType> class KernelData {
+class KernelData {
 public:
-  KernelData(const size_t volumes, const size_t kernel_size)
-      : centre_index(-1),                            //
-        voxel_count(kernel_size),                    //
-        X(MatrixType::Zero(volumes, kernel_size)) {} //
-  size_t centre_index;
-  size_t voxel_count;
-  MatrixType X;
+  KernelData() : centre_index(-1) {}
+  KernelData(const ssize_t i) : centre_index(i) {}
+  ssize_t centre_index;
+  std::vector<voxel_type> voxels;
 };
 
-template <class MatrixType> class KernelBase {
+class KernelBase {
 public:
-  KernelBase() : pos({-1, -1, -1}) {}
-  KernelBase(const KernelBase &) : pos({-1, -1, -1}) {}
+  KernelBase(const Header &H) : H(H) {}
+  KernelBase(const KernelBase &) = default;
   // This is just for pre-allocating matrices
   virtual ssize_t estimated_size() const = 0;
+  // This is the interface that kernels must provide
+  virtual KernelData operator()(const voxel_type &) const = 0;
 
 protected:
-  // Store / restore position of image before / after data loading
-  std::array<ssize_t, 3> pos;
-  template <class ImageType> void stash_pos(const ImageType &image) {
-    for (size_t axis = 0; axis != 3; ++axis)
-      pos[axis] = image.index(axis);
-  }
-  template <class ImageType> void restore_pos(ImageType &image) {
-    for (size_t axis = 0; axis != 3; ++axis)
-      image.index(axis) = pos[axis];
-  }
+  const Header H;
 };
 
-template <class MatrixType> class KernelCube : public KernelBase<MatrixType> {
+class KernelCube : public KernelBase {
 public:
-  KernelCube(const std::vector<uint32_t> &extent)
-      : half_extent({int(extent[0] / 2), int(extent[1] / 2), int(extent[2] / 2)}) {
+  KernelCube(const Header &header, const std::vector<uint32_t> &extent)
+      : KernelBase(header),
+        half_extent({int(extent[0] / 2), int(extent[1] / 2), int(extent[2] / 2)}),
+        size(extent[0] * extent[1] * extent[2]),
+        centre_index(size / 2) {
     for (auto e : extent) {
       if (!(e % 2))
         throw Exception("Size of cubic kernel must be an odd integer");
     }
   }
   KernelCube(const KernelCube &) = default;
-  template <class ImageType> void operator()(ImageType &image, KernelData<MatrixType> &data) {
-    assert(data.X.cols() == size());
-    KernelBase<MatrixType>::stash_pos(image);
-    size_t k = 0;
-    for (int z = -half_extent[2]; z <= half_extent[2]; z++) {
-      image.index(2) = wrapindex(z, 2, image.size(2));
-      for (int y = -half_extent[1]; y <= half_extent[1]; y++) {
-        image.index(1) = wrapindex(y, 1, image.size(1));
-        for (int x = -half_extent[0]; x <= half_extent[0]; x++, k++) {
-          image.index(0) = wrapindex(x, 0, image.size(0));
-          data.X.col(k) = image.row(3);
+  KernelData operator()(const voxel_type &pos) const override {
+    KernelData result(centre_index);
+    voxel_type voxel;
+    voxel_type offset;
+    for (offset[2] = -half_extent[2]; offset[2] <= half_extent[2]; ++offset[2]) {
+      voxel[2] = wrapindex(pos[2], offset[2], half_extent[2], H.size(2));
+      for (offset[1] = -half_extent[1]; offset[1] <= half_extent[1]; ++offset[1]) {
+        voxel[1] = wrapindex(pos[1], offset[1], half_extent[1], H.size(1));
+        for (offset[0] = -half_extent[0]; offset[0] <= half_extent[0]; ++offset[0]) {
+          voxel[0] = wrapindex(pos[0], offset[0], half_extent[0], H.size(0));
+          result.voxels.push_back(pos);
         }
       }
     }
-    KernelBase<MatrixType>::restore_pos(image);
-    data.voxel_count = size();
-    data.centre_index = size() / 2;
-  }
-  ssize_t size() const { return estimated_size(); }
-  ssize_t estimated_size() const override {
-    return (2 * half_extent[0] + 1) * (2 * half_extent[1] + 1) * (2 * half_extent[2] + 1);
+    return result;
   }
+  ssize_t estimated_size() const override { return size; }
 
 private:
+  const std::vector<int> dimensions;
   const std::vector<int> half_extent;
+  const ssize_t size;
+  const ssize_t centre_index;
 
   // patch handling at image edges
-  inline size_t wrapindex(int r, int axis, int max) const {
-    int rr = KernelBase<MatrixType>::pos[axis] + r;
+  inline size_t wrapindex(int p, int r, int e, int max) const {
+    int rr = p + r;
     if (rr < 0)
-      rr = half_extent[axis] - r;
+      rr = e - r;
     if (rr >= max)
-      rr = (max - 1) - half_extent[axis] - r;
+      rr = (max - 1) - e - r;
     return rr;
   }
 };
 
-template <class MatrixType> class KernelSphereBase : public KernelBase<MatrixType> {
+class KernelSphereBase : public KernelBase {
 public:
   KernelSphereBase(const Header &voxel_grid, const default_type max_radius)
-      : shared(new Shared(voxel_grid, max_radius)) {}
+      : KernelBase(voxel_grid), shared(new Shared(voxel_grid, max_radius)) {}
 
 protected:
   class Shared {
@@ -296,47 +288,36 @@ template <class MatrixType> class KernelSphereBase : public KernelBase<MatrixTyp
   std::shared_ptr<Shared> shared;
 };
 
-template <class MatrixType> class KernelSphereRatio : public KernelSphereBase<MatrixType> {
+class KernelSphereRatio : public KernelSphereBase {
 public:
   KernelSphereRatio(const Header &voxel_grid, const default_type min_ratio)
-      : KernelSphereBase<MatrixType>(voxel_grid, compute_max_radius(voxel_grid, min_ratio)),
+      : KernelSphereBase(voxel_grid, compute_max_radius(voxel_grid, min_ratio)),
         min_size(std::ceil(voxel_grid.size(3) * min_ratio)) {}
-  template <class ImageType> void operator()(ImageType &image, KernelData<MatrixType> &data) {
-    KernelBase<MatrixType>::stash_pos(image);
-    data.voxel_count = 0;
+  KernelData operator()(const voxel_type &pos) const override {
+    KernelData result(0);
     default_type prev_distance = -std::numeric_limits<default_type>::infinity();
-    auto map_it = KernelSphereBase<MatrixType>::shared->begin();
-    while (map_it != KernelSphereBase<MatrixType>::shared->end()) {
+    auto map_it = shared->begin();
+    while (map_it != shared->end()) {
       // If there's a tie in distances, want to include all such offsets in the kernel,
       //   even if the size of the utilised kernel extends beyond the minimum size
-      if (map_it->first != prev_distance && data.voxel_count >= min_size)
+      if (map_it->first != prev_distance && result.voxels.size() >= min_size)
         break;
-      for (size_t axis = 0; axis != 3; ++axis)
-        image.index(axis) = KernelBase<MatrixType>::pos[axis] + map_it->second[axis];
-      if (!is_out_of_bounds(image, 0, 3)) {
-        // Is this larger than any kernel this thread has previously encountered?
-        // If so, try to project what the final size is going to be,
-        //   based on the set of voxels with identical distance to this one
-        //   all getting included in the kernel
-        if (data.voxel_count == data.X.cols()) {
-          size_t extra_cols = 1;
-          auto forward_search = map_it;
-          for (++forward_search;
-               forward_search != KernelSphereBase<MatrixType>::shared->end() && forward_search->first == map_it->first;
-               ++forward_search)
-            ++extra_cols;
-          data.X.conservativeResize(data.X.rows(), data.voxel_count + extra_cols);
-        }
-        data.X.col(data.voxel_count) = image.row(3);
-        prev_distance = map_it->first;
-        ++data.voxel_count;
-      }
+      const voxel_type voxel({pos[0] + map_it->second[0],   //
+                              pos[1] + map_it->second[1],   //
+                              pos[2] + map_it->second[2]}); //
+      if (!is_out_of_bounds(H, voxel, 0, 3))
+        result.voxels.push_back(voxel);
+      prev_distance = map_it->first;
       ++map_it;
     }
-    if (map_it == KernelSphereBase<MatrixType>::shared->end())
-      throw Exception("Inadequate spherical kernel initialisation");
-    KernelBase<MatrixType>::restore_pos(image);
-    data.centre_index = 0;
+    if (map_it == shared->end()) {
+      throw Exception(                                                                   //
+          "Inadequate spherical kernel initialisation "                                  //
+          + "(lookup table " + str(std::distance(shared->begin(), shared->end())) + "; " //
+          + "min size " + str(min_size) + "; "                                           //
+          + "read size " + str(result.voxels.size()) + ")");                             //
+    }
+    return result;
   }
   ssize_t estimated_size() const override { return min_size; }
 
@@ -359,32 +340,24 @@ template <class MatrixType> class KernelSphereRatio : public KernelSphereBase<Ma
   }
 };
 
-template <class MatrixType> class KernelSphereFixedRadius : public KernelSphereBase<MatrixType> {
+class KernelSphereFixedRadius : public KernelSphereBase {
 public:
   KernelSphereFixedRadius(const Header &voxel_grid, const default_type radius)
-      : KernelSphereBase<MatrixType>(voxel_grid, radius),
-        maximum_size(std::distance(KernelSphereBase<MatrixType>::shared->begin(),  //
-                                   KernelSphereBase<MatrixType>::shared->end())) { //
+      : KernelSphereBase(voxel_grid, radius),                         //
+        maximum_size(std::distance(shared->begin(), shared->end())) { //
     INFO("Maximum number of voxels in " + str(radius) + "mm fixed-radius kernel is " + str(maximum_size));
   }
-  template <class ImageType> void operator()(ImageType &image, KernelData<MatrixType> &data) {
-    KernelBase<MatrixType>::stash_pos(image);
-    data.voxel_count = 0;
-    default_type prev_distance = -std::numeric_limits<default_type>::infinity();
-    for (auto map_it = KernelSphereBase<MatrixType>::shared->begin();
-         map_it != KernelSphereBase<MatrixType>::shared->end();
-         ++map_it) {
-      for (size_t axis = 0; axis != 3; ++axis)
-        image.index(axis) = KernelBase<MatrixType>::pos[axis] + map_it->second[axis];
-      if (!is_out_of_bounds(image, 0, 3)) {
-        // We should not need to do any matrix size checking here;
-        //   it should have already been allocated to the maximum size of the kernel
-        data.X.col(data.voxel_count) = image.row(3);
-        ++data.voxel_count;
-      }
+  KernelData operator()(const voxel_type &pos) const {
+    KernelData result(0);
+    result.voxels.reserve(maximum_size);
+    for (auto map_it = shared->begin(); map_it != shared->end(); ++map_it) {
+      const voxel_type voxel({pos[0] + map_it->second[0],   //
+                              pos[1] + map_it->second[1],   //
+                              pos[2] + map_it->second[2]}); //
+      if (!is_out_of_bounds(H, voxel, 0, 3))
+        result.voxels.push_back(voxel);
     }
-    KernelBase<MatrixType>::restore_pos(image);
-    data.centre_index = 0;
+    return result;
   }
   ssize_t estimated_size() const override { return maximum_size; }
 
@@ -392,26 +365,26 @@ template <class MatrixType> class KernelSphereFixedRadius : public KernelSphereB
   const ssize_t maximum_size;
 };
 
-template <typename F, class KernelType> class DenoisingFunctor {
+template <typename F> class DenoisingFunctor {
 
 public:
   using MatrixType = Eigen::Matrix<F, Eigen::Dynamic, Eigen::Dynamic>;
   using SValsType = Eigen::VectorXd;
 
   DenoisingFunctor(int ndwi,
-                   KernelType &kernel,
+                   std::shared_ptr<KernelBase> kernel,
                    Image<bool> &mask,
                    Image<real_type> &noise,
                    Image<uint16_t> &rank,
                    Image<uint16_t> &voxels,
                    bool exp1)
-      : data(ndwi, kernel.estimated_size()),
-        kernel(kernel),
+      : kernel(kernel),
         m(ndwi),
         exp1(exp1),
-        XtX(std::min(m, kernel.estimated_size()), std::min(m, kernel.estimated_size())),
-        eig(std::min(m, kernel.estimated_size())),
-        s(std::min(m, kernel.estimated_size())),
+        X(ndwi, kernel->estimated_size()),
+        XtX(std::min(m, kernel->estimated_size()), std::min(m, kernel->estimated_size())),
+        eig(std::min(m, kernel->estimated_size())),
+        s(std::min(m, kernel->estimated_size())),
         mask(mask),
         noise(noise),
         rankmap(rank),
@@ -425,15 +398,19 @@ template <typename F, class KernelType> class DenoisingFunctor {
         return;
     }
 
-    // Load data in local window
-    kernel(dwi, data);
-    auto X = data.X.leftCols(data.voxel_count);
-
-    const ssize_t n = data.voxel_count;
+    // Load list of voxels from which to load data
+    KernelData neighbourhood = (*kernel)({int(dwi.index(0)), int(dwi.index(1)), int(dwi.index(2))});
+    const ssize_t n = neighbourhood.voxels.size();
     const ssize_t r = std::min(m, n);
     const ssize_t q = std::max(m, n);
 
-    if (r > XtX.rows()) {
+    // Expand local storage if necessary
+    if (n > X.cols()) {
+      DEBUG("Expanding data matrix storage from " + str(m) + "x" + str(X.cols()) + " to " + str(m) + "x" + str(n));
+      X.resize(m, n);
+    }
+    if (r > XtX.cols()) {
+      DEBUG("Expanding decomposition matrix storage from " + str(X.rows()) + " to " + str(r));
       XtX.resize(r, r);
       s.resize(r);
     }
@@ -443,15 +420,18 @@ template <typename F, class KernelType> class DenoisingFunctor {
     //   due to use of block oberations to prevent memory re-allocation
     //   in the presence of variation in kernel sizes
 #ifndef NDEBUG
+    X.fill(std::numeric_limits<F>::signaling_NaN());
     XtX.fill(std::numeric_limits<F>::signaling_NaN());
     s.fill(std::numeric_limits<default_type>::signaling_NaN());
 #endif
 
+    load_data(dwi, neighbourhood.voxels);
+
     // Compute Eigendecomposition:
     if (m <= n)
-      XtX.topLeftCorner(r, r).template triangularView<Eigen::Lower>() = X * X.adjoint();
+      XtX.topLeftCorner(r, r).template triangularView<Eigen::Lower>() = X.leftCols(n) * X.leftCols(n).adjoint();
     else
-      XtX.topLeftCorner(r, r).template triangularView<Eigen::Lower>() = X.adjoint() * X;
+      XtX.topLeftCorner(r, r).template triangularView<Eigen::Lower>() = X.leftCols(n).adjoint() * X.leftCols(n);
     eig.compute(XtX.topLeftCorner(r, r));
     // eigenvalues sorted in increasing order:
     s.head(r) = eig.eigenvalues().template cast<double>();
@@ -480,16 +460,18 @@ template <typename F, class KernelType> class DenoisingFunctor {
       s.head(cutoff_p).setZero();
       s.segment(cutoff_p, r - cutoff_p).setOnes();
       if (m <= n)
-        X.col(data.centre_index) = eig.eigenvectors() * (s.head(r).cast<F>().asDiagonal() *
-                                                         (eig.eigenvectors().adjoint() * X.col(data.centre_index)));
+        X.col(neighbourhood.centre_index) =
+            eig.eigenvectors() *
+            (s.head(r).cast<F>().asDiagonal() * (eig.eigenvectors().adjoint() * X.col(neighbourhood.centre_index)));
       else
-        X.col(data.centre_index) = X * (eig.eigenvectors() * (s.head(r).cast<F>().asDiagonal() *
-                                                              eig.eigenvectors().adjoint().col(data.centre_index)));
+        X.col(neighbourhood.centre_index) =
+            X.leftCols(n) * (eig.eigenvectors() * (s.head(r).cast<F>().asDiagonal() *
+                                                   eig.eigenvectors().adjoint().col(neighbourhood.centre_index)));
     }
 
     // Store output
     assign_pos_of(dwi).to(out);
-    out.row(3) = X.col(data.centre_index);
+    out.row(3) = X.col(neighbourhood.centre_index);
 
     // store noise map if requested:
     if (noise.valid()) {
@@ -509,10 +491,10 @@ template <typename F, class KernelType> class DenoisingFunctor {
   }
 
 private:
-  KernelData<MatrixType> data;
-  KernelType kernel;
+  std::shared_ptr<KernelBase> kernel;
   const ssize_t m;
   const bool exp1;
+  MatrixType X;
   MatrixType XtX;
   Eigen::SelfAdjointEigenSolver<MatrixType> eig;
   SValsType s;
@@ -521,16 +503,25 @@ template <typename F, class KernelType> class DenoisingFunctor {
   Image<real_type> noise;
   Image<uint16_t> rankmap;
   Image<uint16_t> voxelsmap;
+
+  template <typename ImageType> void load_data(ImageType &image, const std::vector<voxel_type> &voxels) {
+    const voxel_type pos({int(image.index(0)), int(image.index(1)), int(image.index(2))});
+    for (ssize_t i = 0; i != voxels.size(); ++i) {
+      assign_pos_of(voxels[i], 0, 3).to(image);
+      X.col(i) = image.row(3);
+    }
+    assign_pos_of(pos, 0, 3).to(image);
+  }
 };
 
-template <typename T, class KernelType>
+template <typename T>
 void run(Header &data,
          Image<bool> &mask,
          Image<real_type> &noise,
          Image<uint16_t> &rank,
          Image<uint16_t> &voxels,
          const std::string &output_name,
-         KernelType &kernel,
+         std::shared_ptr<KernelBase> kernel,
          bool exp1) {
   auto input = data.get_image<T>().with_direct_io(3);
   // create output
@@ -538,81 +529,10 @@ void run(Header &data,
   header.datatype() = DataType::from<T>();
   auto output = Image<T>::create(output_name, header);
   // run
-  DenoisingFunctor<T, KernelType> func(data.size(3), kernel, mask, noise, rank, voxels, exp1);
+  DenoisingFunctor<T> func(data.size(3), kernel, mask, noise, rank, voxels, exp1);
   ThreadedLoop("running MP-PCA denoising", data, 0, 3).run(func, input, output);
 }
 
-template <typename T>
-void run(Header &data,
-         Image<bool> &mask,
-         Image<real_type> &noise,
-         Image<uint16_t> &rank,
-         Image<uint16_t> &voxels,
-         const std::string &output_name,
-         bool exp1) {
-  using MatrixType = Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>;
-
-  auto opt = get_options("shape");
-  const shape_type shape = opt.empty() ? shape_type::SPHERE : shape_type((int)(opt[0][0]));
-
-  switch (shape) {
-  case shape_type::SPHERE: {
-    // TODO Could infer that user wants a cuboid kernel if -extent is used, even if -shape is not
-    if (!get_options("extent").empty())
-      throw Exception("-extent option does not apply to spherical kernel");
-    opt = get_options("radius_mm");
-    if (opt.size()) {
-      KernelSphereFixedRadius<MatrixType> kernel(data, opt[0][0]);
-      run<T, KernelSphereFixedRadius<MatrixType>>(data, mask, noise, rank, voxels, output_name, kernel, exp1);
-      return;
-    }
-    const default_type min_ratio = get_option_value("-radius_ratio", sphere_multiplier_default);
-    KernelSphereRatio<MatrixType> kernel(data, min_ratio);
-    run<T, KernelSphereRatio<MatrixType>>(data, mask, noise, rank, voxels, output_name, kernel, exp1);
-    return;
-  }
-  case shape_type::CUBOID: {
-    if (!get_options("radius_mm").size() || !get_options("radius_ratio").empty())
-      throw Exception("-radius_* options are inapplicable if cuboid kernel shape is selected");
-    opt = get_options("extent");
-    std::vector<uint32_t> extent;
-    if (!opt.empty()) {
-      extent = parse_ints<uint32_t>(opt[0][0]);
-      if (extent.size() == 1)
-        extent = {extent[0], extent[0], extent[0]};
-      if (extent.size() != 3)
-        throw Exception("-extent must be either a scalar or a list of length 3");
-      for (int i = 0; i < 3; i++) {
-        if (!(extent[i] & 1))
-          throw Exception("-extent must be a (list of) odd numbers");
-        if (extent[i] > data.size(i))
-          throw Exception("-extent must not exceed the image dimensions");
-      }
-    } else {
-      uint32_t e = 1;
-      while (Math::pow3(e) < data.size(3))
-        e += 2;
-      extent = {std::min(e, uint32_t(data.size(0))),  //
-                std::min(e, uint32_t(data.size(1))),  //
-                std::min(e, uint32_t(data.size(2)))}; //
-    }
-    INFO("selected patch size: " + str(extent[0]) + " x " + str(extent[1]) + " x " + str(extent[2]) + ".");
-
-    if (std::min<uint32_t>(data.size(3), extent[0] * extent[1] * extent[2]) < 15) {
-      WARN("The number of volumes or the patch size is small. "
-           "This may lead to discretisation effects in the noise level "
-           "and cause inconsistent denoising between adjacent voxels.");
-    }
-
-    KernelCube<MatrixType> kernel(extent);
-    run<T, KernelCube<MatrixType>>(data, mask, noise, rank, voxels, output_name, kernel, exp1);
-    return;
-  }
-  default:
-    assert(false);
-  }
-}
-
 void run() {
   auto dwi = Header::open(argument[0]);
 
@@ -657,25 +577,80 @@ void run() {
     voxels = Image<uint16_t>::create(opt[0][0], header);
   }
 
+  opt = get_options("shape");
+  const shape_type shape = opt.empty() ? shape_type::SPHERE : shape_type((int)(opt[0][0]));
+  std::shared_ptr<KernelBase> kernel;
+
+  switch (shape) {
+  case shape_type::SPHERE: {
+    // TODO Could infer that user wants a cuboid kernel if -extent is used, even if -shape is not
+    if (!get_options("extent").empty())
+      throw Exception("-extent option does not apply to spherical kernel");
+    opt = get_options("radius_mm");
+    if (opt.size())
+      kernel.reset(new KernelSphereFixedRadius(dwi, opt[0][0]));
+    else
+      kernel.reset(new KernelSphereRatio(dwi, get_option_value("-radius_ratio", sphere_multiplier_default)));
+  } break;
+  case shape_type::CUBOID: {
+    if (!get_options("radius_mm").empty() || !get_options("radius_ratio").empty())
+      throw Exception("-radius_* options are inapplicable if cuboid kernel shape is selected");
+    opt = get_options("extent");
+    std::vector<uint32_t> extent;
+    if (!opt.empty()) {
+      extent = parse_ints<uint32_t>(opt[0][0]);
+      if (extent.size() == 1)
+        extent = {extent[0], extent[0], extent[0]};
+      if (extent.size() != 3)
+        throw Exception("-extent must be either a scalar or a list of length 3");
+      for (int i = 0; i < 3; i++) {
+        if (!(extent[i] & 1))
+          throw Exception("-extent must be a (list of) odd numbers");
+        if (extent[i] > dwi.size(i))
+          throw Exception("-extent must not exceed the image dimensions");
+      }
+    } else {
+      uint32_t e = 1;
+      while (Math::pow3(e) < dwi.size(3))
+        e += 2;
+      extent = {std::min(e, uint32_t(dwi.size(0))),  //
+                std::min(e, uint32_t(dwi.size(1))),  //
+                std::min(e, uint32_t(dwi.size(2)))}; //
+    }
+    INFO("selected patch size: " + str(extent[0]) + " x " + str(extent[1]) + " x " + str(extent[2]) + ".");
+
+    if (std::min<uint32_t>(dwi.size(3), extent[0] * extent[1] * extent[2]) < 15) {
+      WARN("The number of volumes or the patch size is small. "
+           "This may lead to discretisation effects in the noise level "
+           "and cause inconsistent denoising between adjacent voxels.");
+    }
+
+    kernel.reset(new KernelCube(dwi, extent));
+  } break;
+  default:
+    assert(false);
+  }
+  assert(kernel);
+
   int prec = get_option_value("datatype", 0); // default: single precision
   if (dwi.datatype().is_complex())
     prec += 2; // support complex input data
   switch (prec) {
   case 0:
     INFO("select real float32 for processing");
-    run<float>(dwi, mask, noise, rank, voxels, argument[1], exp1);
+    run<float>(dwi, mask, noise, rank, voxels, argument[1], kernel, exp1);
     break;
   case 1:
     INFO("select real float64 for processing");
-    run<double>(dwi, mask, noise, rank, voxels, argument[1], exp1);
+    run<double>(dwi, mask, noise, rank, voxels, argument[1], kernel, exp1);
     break;
   case 2:
     INFO("select complex float32 for processing");
-    run<cfloat>(dwi, mask, noise, rank, voxels, argument[1], exp1);
+    run<cfloat>(dwi, mask, noise, rank, voxels, argument[1], kernel, exp1);
     break;
   case 3:
     INFO("select complex float64 for processing");
-    run<cdouble>(dwi, mask, noise, rank, voxels, argument[1], exp1);
+    run<cdouble>(dwi, mask, noise, rank, voxels, argument[1], kernel, exp1);
     break;
   }
 }

From 6ae5583f64f6532e10c440d3b0f141256ebe5c8e Mon Sep 17 00:00:00 2001
From: Robert Smith <robert.smith@florey.edu.au>
Date: Tue, 5 Nov 2024 16:18:55 +1100
Subject: [PATCH 06/34] dwidenoise: Change code handling of estimator selection

---
 cmd/dwidenoise.cpp | 39 +++++++++++++++++++++++++++------------
 1 file changed, 27 insertions(+), 12 deletions(-)

diff --git a/cmd/dwidenoise.cpp b/cmd/dwidenoise.cpp
index 7214150a9d..294d69df2b 100644
--- a/cmd/dwidenoise.cpp
+++ b/cmd/dwidenoise.cpp
@@ -25,6 +25,7 @@ using namespace App;
 
 const std::vector<std::string> dtypes = {"float32", "float64"};
 const std::vector<std::string> estimators = {"exp1", "exp2"};
+enum class estimator_type { EXP1, EXP2 };
 
 const std::vector<std::string> shapes = {"cuboid", "sphere"};
 enum class shape_type { CUBOID, SPHERE };
@@ -312,7 +313,7 @@ class KernelSphereRatio : public KernelSphereBase {
     }
     if (map_it == shared->end()) {
       throw Exception(                                                                   //
-          "Inadequate spherical kernel initialisation "                                  //
+          std::string("Inadequate spherical kernel initialisation ")                     //
           + "(lookup table " + str(std::distance(shared->begin(), shared->end())) + "; " //
           + "min size " + str(min_size) + "; "                                           //
           + "read size " + str(result.voxels.size()) + ")");                             //
@@ -377,10 +378,10 @@ template <typename F> class DenoisingFunctor {
                    Image<real_type> &noise,
                    Image<uint16_t> &rank,
                    Image<uint16_t> &voxels,
-                   bool exp1)
+                   estimator_type estimator)
       : kernel(kernel),
         m(ndwi),
-        exp1(exp1),
+        estimator(estimator),
         X(ndwi, kernel->estimated_size()),
         XtX(std::min(m, kernel->estimated_size()), std::min(m, kernel->estimated_size())),
         eig(std::min(m, kernel->estimated_size())),
@@ -445,7 +446,18 @@ template <typename F> class DenoisingFunctor {
     {                               // (as opposed to the paper where p is defined as the number of signal components)
       double lam = std::max(s[p], 0.0) / q;
       clam += lam;
-      double gam = double(p + 1) / (exp1 ? q : q - (r - p - 1));
+      double denominator;
+      switch (estimator) {
+      case estimator_type::EXP1:
+        denominator = q;
+        break;
+      case estimator_type::EXP2:
+        denominator = q - (r - p - 1);
+        break;
+      default:
+        assert(false);
+      }
+      double gam = double(p + 1) / denominator;
       double sigsq1 = clam / double(p + 1);
       double sigsq2 = (lam - lam_r) / (4.0 * std::sqrt(gam));
       // sigsq2 > sigsq1 if signal else noise
@@ -493,7 +505,7 @@ template <typename F> class DenoisingFunctor {
 private:
   std::shared_ptr<KernelBase> kernel;
   const ssize_t m;
-  const bool exp1;
+  const estimator_type estimator;
   MatrixType X;
   MatrixType XtX;
   Eigen::SelfAdjointEigenSolver<MatrixType> eig;
@@ -522,14 +534,14 @@ void run(Header &data,
          Image<uint16_t> &voxels,
          const std::string &output_name,
          std::shared_ptr<KernelBase> kernel,
-         bool exp1) {
+         estimator_type estimator) {
   auto input = data.get_image<T>().with_direct_io(3);
   // create output
   Header header(data);
   header.datatype() = DataType::from<T>();
   auto output = Image<T>::create(output_name, header);
   // run
-  DenoisingFunctor<T> func(data.size(3), kernel, mask, noise, rank, voxels, exp1);
+  DenoisingFunctor<T> func(data.size(3), kernel, mask, noise, rank, voxels, estimator);
   ThreadedLoop("running MP-PCA denoising", data, 0, 3).run(func, input, output);
 }
 
@@ -546,7 +558,10 @@ void run() {
     check_dimensions(mask, dwi, 0, 3);
   }
 
-  bool exp1 = get_option_value("estimator", 1) == 0; // default: Exp2 (unbiased estimator)
+  estimator_type estimator = estimator_type::EXP2; // default: Exp2 (unbiased estimator)
+  opt = get_options("estimator");
+  if (opt.size())
+    estimator = estimator_type(int(opt[0][0]));
 
   Image<real_type> noise;
   opt = get_options("noise");
@@ -638,19 +653,19 @@ void run() {
   switch (prec) {
   case 0:
     INFO("select real float32 for processing");
-    run<float>(dwi, mask, noise, rank, voxels, argument[1], kernel, exp1);
+    run<float>(dwi, mask, noise, rank, voxels, argument[1], kernel, estimator);
     break;
   case 1:
     INFO("select real float64 for processing");
-    run<double>(dwi, mask, noise, rank, voxels, argument[1], kernel, exp1);
+    run<double>(dwi, mask, noise, rank, voxels, argument[1], kernel, estimator);
     break;
   case 2:
     INFO("select complex float32 for processing");
-    run<cfloat>(dwi, mask, noise, rank, voxels, argument[1], kernel, exp1);
+    run<cfloat>(dwi, mask, noise, rank, voxels, argument[1], kernel, estimator);
     break;
   case 3:
     INFO("select complex float64 for processing");
-    run<cdouble>(dwi, mask, noise, rank, voxels, argument[1], kernel, exp1);
+    run<cdouble>(dwi, mask, noise, rank, voxels, argument[1], kernel, estimator);
     break;
   }
 }

From e3bb26b1b9ca9a24d77ad4cf27042d2ab0790439 Mon Sep 17 00:00:00 2001
From: Robert Smith <robert.smith@florey.edu.au>
Date: Wed, 6 Nov 2024 01:03:55 +1100
Subject: [PATCH 07/34] dwidenoise: Multiple changes to kernel handling

- Add option -max_dist, which exports a map encoding for each voel the maximal distance between that voxel and a voxel included in the PCA decomposition.
- Refactor KernelData to keep track of the distance between each voxel within the kernel and the voxel being processed.
- Bug fix to spherical kernels; kernel was erroneously being applied from the outside inwards rather than from the origin outwards.
- Change type used for data encoding the spherical kernel, which should improve computational performance.
---
 cmd/dwidenoise.cpp | 138 ++++++++++++++++++++++++++++++---------------
 1 file changed, 93 insertions(+), 45 deletions(-)

diff --git a/cmd/dwidenoise.cpp b/cmd/dwidenoise.cpp
index 294d69df2b..33c27e03a5 100644
--- a/cmd/dwidenoise.cpp
+++ b/cmd/dwidenoise.cpp
@@ -121,7 +121,7 @@ void usage() {
   + OptionGroup("Options for exporting additional data regarding PCA behaviour")
   + Option("noise",
            "The output noise map,"
-           " i.e., the estimated noise level 'sigma' in the data."
+           " i.e., the estimated noise level 'sigma' in the data. "
            "Note that on complex input data,"
            " this will be the total noise level across real and imaginary channels,"
            " so a scale factor sqrt(2) applies.")
@@ -129,6 +129,9 @@ void usage() {
   + Option("rank",
            "The selected signal rank of the output denoised image.")
     + Argument("cutoff").type_image_out()
+  + Option("max_dist",
+           "The maximum distance between a voxel and another voxel that was included in the local denoising patch")
+    + Argument("image").type_image_out()
   + Option("voxels",
            "The number of voxels that contributed to the PCA for processing of each voxel")
     + Argument("image").type_image_out()
@@ -182,13 +185,30 @@ void usage() {
 using real_type = float;
 using voxel_type = Eigen::Array<int, 3, 1>;
 
+class KernelVoxel {
+public:
+  KernelVoxel(const voxel_type offset, const default_type sq_distance) : offset(offset), sq_distance(sq_distance) {}
+  KernelVoxel(const KernelVoxel &) = default;
+  KernelVoxel(KernelVoxel &&) = default;
+  KernelVoxel &operator=(const KernelVoxel &that) {
+    offset = that.offset;
+    sq_distance = that.sq_distance;
+    return *this;
+  }
+  bool operator<(const KernelVoxel &that) const { return sq_distance < that.sq_distance; }
+  default_type distance() const { return std::sqrt(sq_distance); }
+  voxel_type offset;
+  default_type sq_distance;
+};
+
 // Class to encode return information from kernel
 class KernelData {
 public:
-  KernelData() : centre_index(-1) {}
-  KernelData(const ssize_t i) : centre_index(i) {}
+  KernelData() : centre_index(-1), max_distance(-std::numeric_limits<default_type>::infinity()) {}
+  KernelData(const ssize_t i) : centre_index(i), max_distance(-std::numeric_limits<default_type>::infinity()) {}
+  std::vector<KernelVoxel> voxels;
   ssize_t centre_index;
-  std::vector<voxel_type> voxels;
+  default_type max_distance;
 };
 
 class KernelBase {
@@ -227,10 +247,15 @@ class KernelCube : public KernelBase {
         voxel[1] = wrapindex(pos[1], offset[1], half_extent[1], H.size(1));
         for (offset[0] = -half_extent[0]; offset[0] <= half_extent[0]; ++offset[0]) {
           voxel[0] = wrapindex(pos[0], offset[0], half_extent[0], H.size(0));
-          result.voxels.push_back(pos);
+          const default_type sq_distance = Math::pow2((pos[0] - voxel[0]) * H.spacing(0)) +
+                                           Math::pow2((pos[1] - voxel[1]) * H.spacing(1)) +
+                                           Math::pow2((pos[2] - voxel[2]) * H.spacing(2));
+          result.voxels.push_back(KernelVoxel(voxel, sq_distance));
+          result.max_distance = std::max(result.max_distance, sq_distance);
         }
       }
     }
+    result.max_distance = std::sqrt(result.max_distance);
     return result;
   }
   ssize_t estimated_size() const override { return size; }
@@ -260,14 +285,15 @@ class KernelSphereBase : public KernelBase {
 protected:
   class Shared {
   public:
-    using MapType = std::multimap<float, std::array<int, 3>>;
+    using TableType = std::vector<KernelVoxel>;
     Shared(const Header &voxel_grid, const default_type max_radius) {
       const default_type max_radius_sq = Math::pow2(max_radius);
-      const std::array<int, 3> half_extents({int(std::ceil(max_radius / voxel_grid.spacing(0))),   //
-                                             int(std::ceil(max_radius / voxel_grid.spacing(1))),   //
-                                             int(std::ceil(max_radius / voxel_grid.spacing(2)))}); //
+      const voxel_type half_extents({int(std::ceil(max_radius / voxel_grid.spacing(0))),   //
+                                     int(std::ceil(max_radius / voxel_grid.spacing(1))),   //
+                                     int(std::ceil(max_radius / voxel_grid.spacing(2)))}); //
       // Build the searchlight
-      std::array<int, 3> offset;
+      data.reserve((2 * half_extents[0] + 1) * (2 * half_extents[1] + 1) * (2 * half_extents[2] + 1));
+      voxel_type offset;
       for (offset[2] = -half_extents[2]; offset[2] <= half_extents[2]; ++offset[2]) {
         for (offset[1] = -half_extents[1]; offset[1] <= half_extents[1]; ++offset[1]) {
           for (offset[0] = -half_extents[0]; offset[0] <= half_extents[0]; ++offset[0]) {
@@ -275,16 +301,17 @@ class KernelSphereBase : public KernelBase {
                                                   + Math::pow2(offset[1] * voxel_grid.spacing(1))  //
                                                   + Math::pow2(offset[2] * voxel_grid.spacing(2)); //
             if (squared_distance <= max_radius_sq)
-              data.insert({squared_distance, offset});
+              data.emplace_back(KernelVoxel(offset, squared_distance));
           }
         }
       }
+      std::sort(data.begin(), data.end());
     }
-    MapType::const_iterator begin() const { return data.begin(); }
-    MapType::const_iterator end() const { return data.end(); }
+    TableType::const_iterator begin() const { return data.begin(); }
+    TableType::const_iterator end() const { return data.end(); }
 
   private:
-    MapType data;
+    TableType data;
   };
   std::shared_ptr<Shared> shared;
 };
@@ -296,28 +323,29 @@ class KernelSphereRatio : public KernelSphereBase {
         min_size(std::ceil(voxel_grid.size(3) * min_ratio)) {}
   KernelData operator()(const voxel_type &pos) const override {
     KernelData result(0);
-    default_type prev_distance = -std::numeric_limits<default_type>::infinity();
-    auto map_it = shared->begin();
-    while (map_it != shared->end()) {
+    auto table_it = shared->begin();
+    while (table_it != shared->end()) {
       // If there's a tie in distances, want to include all such offsets in the kernel,
       //   even if the size of the utilised kernel extends beyond the minimum size
-      if (map_it->first != prev_distance && result.voxels.size() >= min_size)
+      if (result.voxels.size() >= min_size && table_it->sq_distance != result.max_distance)
         break;
-      const voxel_type voxel({pos[0] + map_it->second[0],   //
-                              pos[1] + map_it->second[1],   //
-                              pos[2] + map_it->second[2]}); //
-      if (!is_out_of_bounds(H, voxel, 0, 3))
-        result.voxels.push_back(voxel);
-      prev_distance = map_it->first;
-      ++map_it;
+      const voxel_type voxel({pos[0] + table_it->offset[0],   //
+                              pos[1] + table_it->offset[1],   //
+                              pos[2] + table_it->offset[2]}); //
+      if (!is_out_of_bounds(H, voxel, 0, 3)) {
+        result.voxels.push_back(KernelVoxel(voxel, table_it->sq_distance));
+        result.max_distance = table_it->sq_distance;
+      }
+      ++table_it;
     }
-    if (map_it == shared->end()) {
+    if (table_it == shared->end()) {
       throw Exception(                                                                   //
           std::string("Inadequate spherical kernel initialisation ")                     //
           + "(lookup table " + str(std::distance(shared->begin(), shared->end())) + "; " //
           + "min size " + str(min_size) + "; "                                           //
           + "read size " + str(result.voxels.size()) + ")");                             //
     }
+    result.max_distance = std::sqrt(result.max_distance);
     return result;
   }
   ssize_t estimated_size() const override { return min_size; }
@@ -332,9 +360,9 @@ class KernelSphereRatio : public KernelSphereBase {
     const default_type voxel_volume = voxel_grid.spacing(0) * voxel_grid.spacing(1) * voxel_grid.spacing(2);
     const default_type sphere_volume = 8.0 * num_volumes * min_ratio * voxel_volume;
     const default_type approx_radius = std::sqrt(sphere_volume * 0.75 / Math::pi);
-    const std::array<int, 3> half_extents({int(std::ceil(approx_radius / voxel_grid.spacing(0))),   //
-                                           int(std::ceil(approx_radius / voxel_grid.spacing(1))),   //
-                                           int(std::ceil(approx_radius / voxel_grid.spacing(2)))}); //
+    const voxel_type half_extents({int(std::ceil(approx_radius / voxel_grid.spacing(0))),   //
+                                   int(std::ceil(approx_radius / voxel_grid.spacing(1))),   //
+                                   int(std::ceil(approx_radius / voxel_grid.spacing(2)))}); //
     return std::max({half_extents[0] * voxel_grid.spacing(0),
                      half_extents[1] * voxel_grid.spacing(1),
                      half_extents[2] * voxel_grid.spacing(2)});
@@ -352,12 +380,15 @@ class KernelSphereFixedRadius : public KernelSphereBase {
     KernelData result(0);
     result.voxels.reserve(maximum_size);
     for (auto map_it = shared->begin(); map_it != shared->end(); ++map_it) {
-      const voxel_type voxel({pos[0] + map_it->second[0],   //
-                              pos[1] + map_it->second[1],   //
-                              pos[2] + map_it->second[2]}); //
-      if (!is_out_of_bounds(H, voxel, 0, 3))
-        result.voxels.push_back(voxel);
+      const voxel_type voxel({pos[0] + map_it->offset[0],   //
+                              pos[1] + map_it->offset[1],   //
+                              pos[2] + map_it->offset[2]}); //
+      if (!is_out_of_bounds(H, voxel, 0, 3)) {
+        result.voxels.push_back(KernelVoxel(voxel, map_it->sq_distance));
+        result.max_distance = map_it->sq_distance;
+      }
     }
+    result.max_distance = std::sqrt(result.max_distance);
     return result;
   }
   ssize_t estimated_size() const override { return maximum_size; }
@@ -377,6 +408,7 @@ template <typename F> class DenoisingFunctor {
                    Image<bool> &mask,
                    Image<real_type> &noise,
                    Image<uint16_t> &rank,
+                   Image<float> &max_dist,
                    Image<uint16_t> &voxels,
                    estimator_type estimator)
       : kernel(kernel),
@@ -389,6 +421,7 @@ template <typename F> class DenoisingFunctor {
         mask(mask),
         noise(noise),
         rankmap(rank),
+        maxdistmap(max_dist),
         voxelsmap(voxels) {}
 
   template <typename ImageType> void operator()(ImageType &dwi, ImageType &out) {
@@ -400,7 +433,7 @@ template <typename F> class DenoisingFunctor {
     }
 
     // Load list of voxels from which to load data
-    KernelData neighbourhood = (*kernel)({int(dwi.index(0)), int(dwi.index(1)), int(dwi.index(2))});
+    const KernelData neighbourhood = (*kernel)({int(dwi.index(0)), int(dwi.index(1)), int(dwi.index(2))});
     const ssize_t n = neighbourhood.voxels.size();
     const ssize_t r = std::min(m, n);
     const ssize_t q = std::max(m, n);
@@ -485,17 +518,19 @@ template <typename F> class DenoisingFunctor {
     assign_pos_of(dwi).to(out);
     out.row(3) = X.col(neighbourhood.centre_index);
 
-    // store noise map if requested:
+    // Store additional output maps if requested
     if (noise.valid()) {
       assign_pos_of(dwi, 0, 3).to(noise);
       noise.value() = real_type(std::sqrt(sigma2));
     }
-    // store rank map if requested:
     if (rankmap.valid()) {
       assign_pos_of(dwi, 0, 3).to(rankmap);
       rankmap.value() = uint16_t(r - cutoff_p);
     }
-    // store number of voxels map if requested:
+    if (maxdistmap.valid()) {
+      assign_pos_of(dwi, 0, 3).to(maxdistmap);
+      maxdistmap.value() = neighbourhood.max_distance;
+    }
     if (voxelsmap.valid()) {
       assign_pos_of(dwi, 0, 3).to(voxelsmap);
       voxelsmap.value() = n;
@@ -514,12 +549,13 @@ template <typename F> class DenoisingFunctor {
   Image<bool> mask;
   Image<real_type> noise;
   Image<uint16_t> rankmap;
+  Image<float> maxdistmap;
   Image<uint16_t> voxelsmap;
 
-  template <typename ImageType> void load_data(ImageType &image, const std::vector<voxel_type> &voxels) {
+  template <typename ImageType> void load_data(ImageType &image, const std::vector<KernelVoxel> &voxels) {
     const voxel_type pos({int(image.index(0)), int(image.index(1)), int(image.index(2))});
     for (ssize_t i = 0; i != voxels.size(); ++i) {
-      assign_pos_of(voxels[i], 0, 3).to(image);
+      assign_pos_of(voxels[i].offset, 0, 3).to(image);
       X.col(i) = image.row(3);
     }
     assign_pos_of(pos, 0, 3).to(image);
@@ -531,6 +567,7 @@ void run(Header &data,
          Image<bool> &mask,
          Image<real_type> &noise,
          Image<uint16_t> &rank,
+         Image<float> &max_dist,
          Image<uint16_t> &voxels,
          const std::string &output_name,
          std::shared_ptr<KernelBase> kernel,
@@ -541,7 +578,7 @@ void run(Header &data,
   header.datatype() = DataType::from<T>();
   auto output = Image<T>::create(output_name, header);
   // run
-  DenoisingFunctor<T> func(data.size(3), kernel, mask, noise, rank, voxels, estimator);
+  DenoisingFunctor<T> func(data.size(3), kernel, mask, noise, rank, max_dist, voxels, estimator);
   ThreadedLoop("running MP-PCA denoising", data, 0, 3).run(func, input, output);
 }
 
@@ -582,6 +619,17 @@ void run() {
     rank = Image<uint16_t>::create(opt[0][0], header);
   }
 
+  Image<float> max_dist;
+  opt = get_options("max_dist");
+  if (!opt.empty()) {
+    Header header(dwi);
+    header.ndim() = 3;
+    header.datatype() = DataType::Float32;
+    header.datatype().set_byte_order_native();
+    header.reset_intensity_scaling();
+    max_dist = Image<float>::create(opt[0][0], header);
+  }
+
   Image<uint16_t> voxels;
   opt = get_options("voxels");
   if (!opt.empty()) {
@@ -653,19 +701,19 @@ void run() {
   switch (prec) {
   case 0:
     INFO("select real float32 for processing");
-    run<float>(dwi, mask, noise, rank, voxels, argument[1], kernel, estimator);
+    run<float>(dwi, mask, noise, rank, max_dist, voxels, argument[1], kernel, estimator);
     break;
   case 1:
     INFO("select real float64 for processing");
-    run<double>(dwi, mask, noise, rank, voxels, argument[1], kernel, estimator);
+    run<double>(dwi, mask, noise, rank, max_dist, voxels, argument[1], kernel, estimator);
     break;
   case 2:
     INFO("select complex float32 for processing");
-    run<cfloat>(dwi, mask, noise, rank, voxels, argument[1], kernel, estimator);
+    run<cfloat>(dwi, mask, noise, rank, max_dist, voxels, argument[1], kernel, estimator);
     break;
   case 3:
     INFO("select complex float64 for processing");
-    run<cdouble>(dwi, mask, noise, rank, voxels, argument[1], kernel, estimator);
+    run<cdouble>(dwi, mask, noise, rank, max_dist, voxels, argument[1], kernel, estimator);
     break;
   }
 }

From 750cfd96b1124d163447d777388db5814e3ea97c Mon Sep 17 00:00:00 2001
From: Robert Smith <robert.smith@florey.edu.au>
Date: Wed, 6 Nov 2024 13:42:30 +1100
Subject: [PATCH 08/34] dwidenoise: Cleanups suggested by clang-tidy

---
 cmd/dwidenoise.cpp | 39 ++++++++++++++++++++++++---------------
 1 file changed, 24 insertions(+), 15 deletions(-)

diff --git a/cmd/dwidenoise.cpp b/cmd/dwidenoise.cpp
index 33c27e03a5..1a5d9c5950 100644
--- a/cmd/dwidenoise.cpp
+++ b/cmd/dwidenoise.cpp
@@ -187,14 +187,20 @@ using voxel_type = Eigen::Array<int, 3, 1>;
 
 class KernelVoxel {
 public:
-  KernelVoxel(const voxel_type offset, const default_type sq_distance) : offset(offset), sq_distance(sq_distance) {}
+  KernelVoxel(const voxel_type &offset, const default_type sq_distance) : offset(offset), sq_distance(sq_distance) {}
   KernelVoxel(const KernelVoxel &) = default;
   KernelVoxel(KernelVoxel &&) = default;
+  ~KernelVoxel() {}
   KernelVoxel &operator=(const KernelVoxel &that) {
     offset = that.offset;
     sq_distance = that.sq_distance;
     return *this;
   }
+  KernelVoxel &operator=(KernelVoxel &&that) {
+    offset = that.offset;
+    sq_distance = that.sq_distance;
+    return *this;
+  }
   bool operator<(const KernelVoxel &that) const { return sq_distance < that.sq_distance; }
   default_type distance() const { return std::sqrt(sq_distance); }
   voxel_type offset;
@@ -229,7 +235,7 @@ class KernelCube : public KernelBase {
   KernelCube(const Header &header, const std::vector<uint32_t> &extent)
       : KernelBase(header),
         half_extent({int(extent[0] / 2), int(extent[1] / 2), int(extent[2] / 2)}),
-        size(extent[0] * extent[1] * extent[2]),
+        size(size_t(extent[0]) * size_t(extent[1]) * size_t(extent[2])),
         centre_index(size / 2) {
     for (auto e : extent) {
       if (!(e % 2))
@@ -281,6 +287,7 @@ class KernelSphereBase : public KernelBase {
 public:
   KernelSphereBase(const Header &voxel_grid, const default_type max_radius)
       : KernelBase(voxel_grid), shared(new Shared(voxel_grid, max_radius)) {}
+  virtual ~KernelSphereBase() {}
 
 protected:
   class Shared {
@@ -292,7 +299,7 @@ class KernelSphereBase : public KernelBase {
                                      int(std::ceil(max_radius / voxel_grid.spacing(1))),   //
                                      int(std::ceil(max_radius / voxel_grid.spacing(2)))}); //
       // Build the searchlight
-      data.reserve((2 * half_extents[0] + 1) * (2 * half_extents[1] + 1) * (2 * half_extents[2] + 1));
+      data.reserve(size_t((2 * half_extents[0] + 1) * (2 * half_extents[1] + 1) * (2 * half_extents[2] + 1)));
       voxel_type offset;
       for (offset[2] = -half_extents[2]; offset[2] <= half_extents[2]; ++offset[2]) {
         for (offset[1] = -half_extents[1]; offset[1] <= half_extents[1]; ++offset[1]) {
@@ -321,6 +328,7 @@ class KernelSphereRatio : public KernelSphereBase {
   KernelSphereRatio(const Header &voxel_grid, const default_type min_ratio)
       : KernelSphereBase(voxel_grid, compute_max_radius(voxel_grid, min_ratio)),
         min_size(std::ceil(voxel_grid.size(3) * min_ratio)) {}
+  ~KernelSphereRatio() {}
   KernelData operator()(const voxel_type &pos) const override {
     KernelData result(0);
     auto table_it = shared->begin();
@@ -376,6 +384,7 @@ class KernelSphereFixedRadius : public KernelSphereBase {
         maximum_size(std::distance(shared->begin(), shared->end())) { //
     INFO("Maximum number of voxels in " + str(radius) + "mm fixed-radius kernel is " + str(maximum_size));
   }
+  ~KernelSphereFixedRadius() {}
   KernelData operator()(const voxel_type &pos) const {
     KernelData result(0);
     result.voxels.reserve(maximum_size);
@@ -477,9 +486,9 @@ template <typename F> class DenoisingFunctor {
     ssize_t cutoff_p = 0;
     for (ssize_t p = 0; p < r; ++p) // p+1 is the number of noise components
     {                               // (as opposed to the paper where p is defined as the number of signal components)
-      double lam = std::max(s[p], 0.0) / q;
+      const double lam = std::max(s[p], 0.0) / q;
       clam += lam;
-      double denominator;
+      double denominator = std::numeric_limits<double>::signaling_NaN();
       switch (estimator) {
       case estimator_type::EXP1:
         denominator = q;
@@ -490,9 +499,9 @@ template <typename F> class DenoisingFunctor {
       default:
         assert(false);
       }
-      double gam = double(p + 1) / denominator;
-      double sigsq1 = clam / double(p + 1);
-      double sigsq2 = (lam - lam_r) / (4.0 * std::sqrt(gam));
+      const double gam = double(p + 1) / denominator;
+      const double sigsq1 = clam / double(p + 1);
+      const double sigsq2 = (lam - lam_r) / (4.0 * std::sqrt(gam));
       // sigsq2 > sigsq1 if signal else noise
       if (sigsq2 < sigsq1) {
         sigma2 = sigsq1;
@@ -529,7 +538,7 @@ template <typename F> class DenoisingFunctor {
     }
     if (maxdistmap.valid()) {
       assign_pos_of(dwi, 0, 3).to(maxdistmap);
-      maxdistmap.value() = neighbourhood.max_distance;
+      maxdistmap.value() = float(neighbourhood.max_distance);
     }
     if (voxelsmap.valid()) {
       assign_pos_of(dwi, 0, 3).to(voxelsmap);
@@ -597,7 +606,7 @@ void run() {
 
   estimator_type estimator = estimator_type::EXP2; // default: Exp2 (unbiased estimator)
   opt = get_options("estimator");
-  if (opt.size())
+  if (!opt.empty())
     estimator = estimator_type(int(opt[0][0]));
 
   Image<real_type> noise;
@@ -650,10 +659,10 @@ void run() {
     if (!get_options("extent").empty())
       throw Exception("-extent option does not apply to spherical kernel");
     opt = get_options("radius_mm");
-    if (opt.size())
-      kernel.reset(new KernelSphereFixedRadius(dwi, opt[0][0]));
+    if (opt.empty())
+      kernel = std::make_shared<KernelSphereRatio>(dwi, get_option_value("-radius_ratio", sphere_multiplier_default));
     else
-      kernel.reset(new KernelSphereRatio(dwi, get_option_value("-radius_ratio", sphere_multiplier_default)));
+      kernel = std::make_shared<KernelSphereFixedRadius>(dwi, opt[0][0]);
   } break;
   case shape_type::CUBOID: {
     if (!get_options("radius_mm").empty() || !get_options("radius_ratio").empty())
@@ -667,7 +676,7 @@ void run() {
       if (extent.size() != 3)
         throw Exception("-extent must be either a scalar or a list of length 3");
       for (int i = 0; i < 3; i++) {
-        if (!(extent[i] & 1))
+        if ((extent[i] & 1) == 0)
           throw Exception("-extent must be a (list of) odd numbers");
         if (extent[i] > dwi.size(i))
           throw Exception("-extent must not exceed the image dimensions");
@@ -688,7 +697,7 @@ void run() {
            "and cause inconsistent denoising between adjacent voxels.");
     }
 
-    kernel.reset(new KernelCube(dwi, extent));
+    kernel = std::make_shared<KernelCube>(dwi, extent);
   } break;
   default:
     assert(false);

From 7c8fc048fb8583af75a7fe6bb8b006a9f3224c7e Mon Sep 17 00:00:00 2001
From: Robert Smith <robert.smith@florey.edu.au>
Date: Wed, 6 Nov 2024 15:30:27 +1100
Subject: [PATCH 09/34] dwidenoise: Fix -radius_ratio option

---
 cmd/dwidenoise.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cmd/dwidenoise.cpp b/cmd/dwidenoise.cpp
index 1a5d9c5950..a4b5add2ac 100644
--- a/cmd/dwidenoise.cpp
+++ b/cmd/dwidenoise.cpp
@@ -660,7 +660,7 @@ void run() {
       throw Exception("-extent option does not apply to spherical kernel");
     opt = get_options("radius_mm");
     if (opt.empty())
-      kernel = std::make_shared<KernelSphereRatio>(dwi, get_option_value("-radius_ratio", sphere_multiplier_default));
+      kernel = std::make_shared<KernelSphereRatio>(dwi, get_option_value("radius_ratio", sphere_multiplier_default));
     else
       kernel = std::make_shared<KernelSphereFixedRadius>(dwi, opt[0][0]);
   } break;

From 239e9940d45405efefca298ed48a64abe5276a34 Mon Sep 17 00:00:00 2001
From: Robert Smith <robert.smith@florey.edu.au>
Date: Thu, 7 Nov 2024 12:42:14 +1100
Subject: [PATCH 10/34] dwidenoise: Further fixes for clang-tidy

---
 cmd/dwidenoise.cpp | 19 ++++++++++++-------
 1 file changed, 12 insertions(+), 7 deletions(-)

diff --git a/cmd/dwidenoise.cpp b/cmd/dwidenoise.cpp
index a4b5add2ac..cd8e009ca7 100644
--- a/cmd/dwidenoise.cpp
+++ b/cmd/dwidenoise.cpp
@@ -196,7 +196,7 @@ class KernelVoxel {
     sq_distance = that.sq_distance;
     return *this;
   }
-  KernelVoxel &operator=(KernelVoxel &&that) {
+  KernelVoxel &operator=(KernelVoxel &&that) noexcept {
     offset = that.offset;
     sq_distance = that.sq_distance;
     return *this;
@@ -221,6 +221,7 @@ class KernelBase {
 public:
   KernelBase(const Header &H) : H(H) {}
   KernelBase(const KernelBase &) = default;
+  virtual ~KernelBase() = default;
   // This is just for pre-allocating matrices
   virtual ssize_t estimated_size() const = 0;
   // This is the interface that kernels must provide
@@ -235,7 +236,7 @@ class KernelCube : public KernelBase {
   KernelCube(const Header &header, const std::vector<uint32_t> &extent)
       : KernelBase(header),
         half_extent({int(extent[0] / 2), int(extent[1] / 2), int(extent[2] / 2)}),
-        size(size_t(extent[0]) * size_t(extent[1]) * size_t(extent[2])),
+        size(ssize_t(extent[0]) * ssize_t(extent[1]) * ssize_t(extent[2])),
         centre_index(size / 2) {
     for (auto e : extent) {
       if (!(e % 2))
@@ -243,6 +244,7 @@ class KernelCube : public KernelBase {
     }
   }
   KernelCube(const KernelCube &) = default;
+  ~KernelCube() final = default;
   KernelData operator()(const voxel_type &pos) const override {
     KernelData result(centre_index);
     voxel_type voxel;
@@ -287,7 +289,8 @@ class KernelSphereBase : public KernelBase {
 public:
   KernelSphereBase(const Header &voxel_grid, const default_type max_radius)
       : KernelBase(voxel_grid), shared(new Shared(voxel_grid, max_radius)) {}
-  virtual ~KernelSphereBase() {}
+  KernelSphereBase(const KernelSphereBase &) = default;
+  virtual ~KernelSphereBase() override {}
 
 protected:
   class Shared {
@@ -299,8 +302,8 @@ class KernelSphereBase : public KernelBase {
                                      int(std::ceil(max_radius / voxel_grid.spacing(1))),   //
                                      int(std::ceil(max_radius / voxel_grid.spacing(2)))}); //
       // Build the searchlight
-      data.reserve(size_t((2 * half_extents[0] + 1) * (2 * half_extents[1] + 1) * (2 * half_extents[2] + 1)));
-      voxel_type offset;
+      data.reserve(size_t(2 * half_extents[0] + 1) * size_t(2 * half_extents[1] + 1) * size_t(2 * half_extents[2] + 1));
+      voxel_type offset({-1, -1, -1});
       for (offset[2] = -half_extents[2]; offset[2] <= half_extents[2]; ++offset[2]) {
         for (offset[1] = -half_extents[1]; offset[1] <= half_extents[1]; ++offset[1]) {
           for (offset[0] = -half_extents[0]; offset[0] <= half_extents[0]; ++offset[0]) {
@@ -328,7 +331,8 @@ class KernelSphereRatio : public KernelSphereBase {
   KernelSphereRatio(const Header &voxel_grid, const default_type min_ratio)
       : KernelSphereBase(voxel_grid, compute_max_radius(voxel_grid, min_ratio)),
         min_size(std::ceil(voxel_grid.size(3) * min_ratio)) {}
-  ~KernelSphereRatio() {}
+  KernelSphereRatio(const KernelSphereRatio &) = default;
+  ~KernelSphereRatio() final = default;
   KernelData operator()(const voxel_type &pos) const override {
     KernelData result(0);
     auto table_it = shared->begin();
@@ -384,7 +388,8 @@ class KernelSphereFixedRadius : public KernelSphereBase {
         maximum_size(std::distance(shared->begin(), shared->end())) { //
     INFO("Maximum number of voxels in " + str(radius) + "mm fixed-radius kernel is " + str(maximum_size));
   }
-  ~KernelSphereFixedRadius() {}
+  KernelSphereFixedRadius(const KernelSphereFixedRadius &) = default;
+  ~KernelSphereFixedRadius() final = default;
   KernelData operator()(const voxel_type &pos) const {
     KernelData result(0);
     result.voxels.reserve(maximum_size);

From b66bf235d0e856bc5ff5364996af5ca5c142afa4 Mon Sep 17 00:00:00 2001
From: Robert Smith <robert.smith@florey.edu.au>
Date: Thu, 7 Nov 2024 12:54:27 +1100
Subject: [PATCH 11/34] dwidenoise: Move noise estimate from member to functor
 scope

---
 cmd/dwidenoise.cpp | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/cmd/dwidenoise.cpp b/cmd/dwidenoise.cpp
index cd8e009ca7..a279436669 100644
--- a/cmd/dwidenoise.cpp
+++ b/cmd/dwidenoise.cpp
@@ -487,7 +487,7 @@ template <typename F> class DenoisingFunctor {
     // Marchenko-Pastur optimal threshold
     const double lam_r = std::max(s[0], 0.0) / q;
     double clam = 0.0;
-    sigma2 = 0.0;
+    double sigma2 = 0.0;
     ssize_t cutoff_p = 0;
     for (ssize_t p = 0; p < r; ++p) // p+1 is the number of noise components
     {                               // (as opposed to the paper where p is defined as the number of signal components)
@@ -559,7 +559,6 @@ template <typename F> class DenoisingFunctor {
   MatrixType XtX;
   Eigen::SelfAdjointEigenSolver<MatrixType> eig;
   SValsType s;
-  double sigma2;
   Image<bool> mask;
   Image<real_type> noise;
   Image<uint16_t> rankmap;

From 41652766689ff9a75db5423e181d9192154f825e Mon Sep 17 00:00:00 2001
From: Robert Smith <robert.smith@florey.edu.au>
Date: Thu, 7 Nov 2024 17:09:45 +1100
Subject: [PATCH 12/34] dwidenoise: Change default spherical kernel size

Change from recollection-from-memory ratio of 1.1 to that reported in the corresponding publication of 1.0/0.85.
---
 cmd/dwidenoise.cpp | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/cmd/dwidenoise.cpp b/cmd/dwidenoise.cpp
index a279436669..a4e97ab089 100644
--- a/cmd/dwidenoise.cpp
+++ b/cmd/dwidenoise.cpp
@@ -29,7 +29,7 @@ enum class estimator_type { EXP1, EXP2 };
 
 const std::vector<std::string> shapes = {"cuboid", "sphere"};
 enum class shape_type { CUBOID, SPHERE };
-constexpr default_type sphere_multiplier_default = 1.1;
+constexpr default_type sphere_multiplier_default = 1.0 / 0.85;
 
 // clang-format off
 void usage() {
@@ -145,8 +145,8 @@ void usage() {
            "Set an absolute spherical kernel radius in mm")
     + Argument("value").type_float(0.0)
   + Option("radius_ratio",
-           "Set the spherical kernel radius as a ratio of number of input volumes "
-           "(default: 1.1)")
+           "Set the spherical kernel size as a ratio of number of voxels to number of input volumes "
+           "(default: ~1.18)")
     + Argument("value").type_float(0.0)
   + Option("extent",
            "Set the patch size of the cuboid filter; "

From aec1d06bb44a85ca15076c25d9d0cadbae302a2f Mon Sep 17 00:00:00 2001
From: Robert Smith <robert.smith@florey.edu.au>
Date: Fri, 8 Nov 2024 13:12:19 +1100
Subject: [PATCH 13/34] dwidenoise: Optimal shrinkage

- Default behaviour is now to use optimal shrinkage based on minimisation of the Frobenius norm.
- Prior behaviour can be accessed using "-filter truncate".
Closes #3022.
---
 cmd/dwidenoise.cpp | 138 ++++++++++++++++++++++++++++++++++++---------
 1 file changed, 110 insertions(+), 28 deletions(-)

diff --git a/cmd/dwidenoise.cpp b/cmd/dwidenoise.cpp
index a4e97ab089..493fae334f 100644
--- a/cmd/dwidenoise.cpp
+++ b/cmd/dwidenoise.cpp
@@ -31,6 +31,9 @@ const std::vector<std::string> shapes = {"cuboid", "sphere"};
 enum class shape_type { CUBOID, SPHERE };
 constexpr default_type sphere_multiplier_default = 1.0 / 0.85;
 
+const std::vector<std::string> filters = {"truncate", "frobenius"};
+enum class filter_type { TRUNCATE, FROBENIUS };
+
 // clang-format off
 void usage() {
 
@@ -76,11 +79,18 @@ void usage() {
     "the command will select the smallest isotropic patch size "
     "that exceeds the number of DW images in the input data; "
     "e.g., 5x5x5 for data with <= 125 DWI volumes, "
-    "7x7x7 for data with <= 343 DWI volumes, etc.";
+    "7x7x7 for data with <= 343 DWI volumes, etc."
+
+  + "By default, optimal value shrinkage based on minimisation of the Frobenius norm "
+    "will be used to attenuate eigenvectors based on the estimated noise level. "
+    "Hard truncation of sub-threshold components"
+    "---which was the behaviour of the dwidenoise command in version 3.0.x---"
+    "can be activated using -filter truncate.";
 
   AUTHOR = "Daan Christiaens (daan.christiaens@kcl.ac.uk)"
            " and Jelle Veraart (jelle.veraart@nyumc.org)"
-           " and J-Donald Tournier (jdtournier@gmail.com)";
+           " and J-Donald Tournier (jdtournier@gmail.com)"
+           " and Robert E. Smith (robert.smith@florey.edu.au)";
 
   REFERENCES
   + "Veraart, J.; Novikov, D.S.; Christiaens, D.; Ades-aron, B.; Sijbers, J. & Fieremans, E. " // Internal
@@ -117,6 +127,10 @@ void usage() {
            "* Exp1: the original estimator used in Veraart et al. (2016), or \n"
            "* Exp2: the improved estimator introduced in Cordero-Grande et al. (2019).")
     + Argument("Exp1/Exp2").type_choice(estimators)
+  + Option("filter",
+           "Modulate how components are filtered based on their eigenvalues; "
+           "options are: " + join(filters, ",") + "; default: frobenius")
+    + Argument("choice").type_choice(filters)
 
   + OptionGroup("Options for exporting additional data regarding PCA behaviour")
   + Option("noise",
@@ -125,10 +139,13 @@ void usage() {
            "Note that on complex input data,"
            " this will be the total noise level across real and imaginary channels,"
            " so a scale factor sqrt(2) applies.")
-    + Argument("level").type_image_out()
+    + Argument("image").type_image_out()
   + Option("rank",
            "The selected signal rank of the output denoised image.")
-    + Argument("cutoff").type_image_out()
+    + Argument("image").type_image_out()
+  + Option("sumweights",
+           "the sum of eigenvector weights contributed to the output image")
+    + Argument("image").type_image_out()
   + Option("max_dist",
            "The maximum distance between a voxel and another voxel that was included in the local denoising patch")
     + Argument("image").type_image_out()
@@ -146,10 +163,11 @@ void usage() {
     + Argument("value").type_float(0.0)
   + Option("radius_ratio",
            "Set the spherical kernel size as a ratio of number of voxels to number of input volumes "
-           "(default: ~1.18)")
+           "(default: 1.0/0.85 ~= 1.18)")
     + Argument("value").type_float(0.0)
+  // TODO Command-line option that allows user to specify minimum absolute number of voxels in kernel
   + Option("extent",
-           "Set the patch size of the cuboid filter; "
+           "Set the patch size of the cuboid kernel; "
            "can be either a single odd integer or a comma-separated triplet of odd integers")
     + Argument("window").type_sequence_int();
 
@@ -184,6 +202,7 @@ void usage() {
 
 using real_type = float;
 using voxel_type = Eigen::Array<int, 3, 1>;
+using vector_type = Eigen::VectorXd;
 
 class KernelVoxel {
 public:
@@ -415,26 +434,31 @@ template <typename F> class DenoisingFunctor {
 
 public:
   using MatrixType = Eigen::Matrix<F, Eigen::Dynamic, Eigen::Dynamic>;
-  using SValsType = Eigen::VectorXd;
 
   DenoisingFunctor(int ndwi,
                    std::shared_ptr<KernelBase> kernel,
+                   filter_type filter,
                    Image<bool> &mask,
                    Image<real_type> &noise,
                    Image<uint16_t> &rank,
+                   Image<float> &sum_weights,
                    Image<float> &max_dist,
                    Image<uint16_t> &voxels,
                    estimator_type estimator)
       : kernel(kernel),
+        filter(filter),
         m(ndwi),
         estimator(estimator),
         X(ndwi, kernel->estimated_size()),
         XtX(std::min(m, kernel->estimated_size()), std::min(m, kernel->estimated_size())),
         eig(std::min(m, kernel->estimated_size())),
         s(std::min(m, kernel->estimated_size())),
+        clam(std::min(m, kernel->estimated_size())),
+        w(std::min(m, kernel->estimated_size())),
         mask(mask),
         noise(noise),
         rankmap(rank),
+        sumweightsmap(sum_weights),
         maxdistmap(max_dist),
         voxelsmap(voxels) {}
 
@@ -461,6 +485,8 @@ template <typename F> class DenoisingFunctor {
       DEBUG("Expanding decomposition matrix storage from " + str(X.rows()) + " to " + str(r));
       XtX.resize(r, r);
       s.resize(r);
+      clam.resize(r);
+      w.resize(r);
     }
 
     // Fill matrices with NaN when in debug mode;
@@ -471,6 +497,8 @@ template <typename F> class DenoisingFunctor {
     X.fill(std::numeric_limits<F>::signaling_NaN());
     XtX.fill(std::numeric_limits<F>::signaling_NaN());
     s.fill(std::numeric_limits<default_type>::signaling_NaN());
+    clam.fill(std::numeric_limits<default_type>::signaling_NaN());
+    w.fill(std::numeric_limits<default_type>::signaling_NaN());
 #endif
 
     load_data(dwi, neighbourhood.voxels);
@@ -486,13 +514,12 @@ template <typename F> class DenoisingFunctor {
 
     // Marchenko-Pastur optimal threshold
     const double lam_r = std::max(s[0], 0.0) / q;
-    double clam = 0.0;
     double sigma2 = 0.0;
     ssize_t cutoff_p = 0;
     for (ssize_t p = 0; p < r; ++p) // p+1 is the number of noise components
     {                               // (as opposed to the paper where p is defined as the number of signal components)
       const double lam = std::max(s[p], 0.0) / q;
-      clam += lam;
+      clam[p] = (p == 0 ? 0.0 : clam[p - 1]) + lam;
       double denominator = std::numeric_limits<double>::signaling_NaN();
       switch (estimator) {
       case estimator_type::EXP1:
@@ -505,7 +532,7 @@ template <typename F> class DenoisingFunctor {
         assert(false);
       }
       const double gam = double(p + 1) / denominator;
-      const double sigsq1 = clam / double(p + 1);
+      const double sigsq1 = clam[p] / double(p + 1);
       const double sigsq2 = (lam - lam_r) / (4.0 * std::sqrt(gam));
       // sigsq2 > sigsq1 if signal else noise
       if (sigsq2 < sigsq1) {
@@ -514,20 +541,38 @@ template <typename F> class DenoisingFunctor {
       }
     }
 
-    if (cutoff_p > 0) {
-      // recombine data using only eigenvectors above threshold:
-      s.head(cutoff_p).setZero();
-      s.segment(cutoff_p, r - cutoff_p).setOnes();
-      if (m <= n)
-        X.col(neighbourhood.centre_index) =
-            eig.eigenvectors() *
-            (s.head(r).cast<F>().asDiagonal() * (eig.eigenvectors().adjoint() * X.col(neighbourhood.centre_index)));
-      else
-        X.col(neighbourhood.centre_index) =
-            X.leftCols(n) * (eig.eigenvectors() * (s.head(r).cast<F>().asDiagonal() *
-                                                   eig.eigenvectors().adjoint().col(neighbourhood.centre_index)));
+    // Generate weights vector
+    double sum_weights = 0.0;
+    switch (filter) {
+    case filter_type::TRUNCATE:
+      w.head(cutoff_p).setZero();
+      w.segment(cutoff_p, r - cutoff_p).setOnes();
+      sum_weights = r - cutoff_p;
+      break;
+    case filter_type::FROBENIUS: {
+      const double beta = r / q;
+      const double threshold = 1.0 + std::sqrt(beta);
+      for (ssize_t i = 0; i != r; ++i) {
+        const double y = clam[i] / (sigma2 * (i + 1));
+        const double nu = y > threshold ? std::sqrt(Math::pow2(Math::pow2(y) - beta - 1.0) - (4.0 * beta)) / y : 0.0;
+        w[i] = nu / y;
+        sum_weights += w[i];
+      }
+    } break;
+    default:
+      assert(false);
     }
 
+    // recombine data using only eigenvectors above threshold:
+    if (m <= n)
+      X.col(neighbourhood.centre_index) =
+          eig.eigenvectors() *
+          (w.head(r).cast<F>().asDiagonal() * (eig.eigenvectors().adjoint() * X.col(neighbourhood.centre_index)));
+    else
+      X.col(neighbourhood.centre_index) =
+          X.leftCols(n) * (eig.eigenvectors() * (w.head(r).cast<F>().asDiagonal() *
+                                                 eig.eigenvectors().adjoint().col(neighbourhood.centre_index)));
+
     // Store output
     assign_pos_of(dwi).to(out);
     out.row(3) = X.col(neighbourhood.centre_index);
@@ -541,6 +586,10 @@ template <typename F> class DenoisingFunctor {
       assign_pos_of(dwi, 0, 3).to(rankmap);
       rankmap.value() = uint16_t(r - cutoff_p);
     }
+    if (sumweightsmap.valid()) {
+      assign_pos_of(dwi, 0, 3).to(sumweightsmap);
+      sumweightsmap.value() = sum_weights;
+    }
     if (maxdistmap.valid()) {
       assign_pos_of(dwi, 0, 3).to(maxdistmap);
       maxdistmap.value() = float(neighbourhood.max_distance);
@@ -552,16 +601,27 @@ template <typename F> class DenoisingFunctor {
   }
 
 private:
+  // Denoising configuration
   std::shared_ptr<KernelBase> kernel;
+  filter_type filter;
   const ssize_t m;
   const estimator_type estimator;
+
+  // Reusable memory
   MatrixType X;
   MatrixType XtX;
   Eigen::SelfAdjointEigenSolver<MatrixType> eig;
-  SValsType s;
+  vector_type s;
+  vector_type clam;
+  vector_type w;
+
   Image<bool> mask;
+
+  // Export images
+  // TODO Group these into a class?
   Image<real_type> noise;
   Image<uint16_t> rankmap;
+  Image<float> sumweightsmap;
   Image<float> maxdistmap;
   Image<uint16_t> voxelsmap;
 
@@ -580,10 +640,12 @@ void run(Header &data,
          Image<bool> &mask,
          Image<real_type> &noise,
          Image<uint16_t> &rank,
+         Image<float> &sum_weights,
          Image<float> &max_dist,
          Image<uint16_t> &voxels,
          const std::string &output_name,
          std::shared_ptr<KernelBase> kernel,
+         filter_type filter,
          estimator_type estimator) {
   auto input = data.get_image<T>().with_direct_io(3);
   // create output
@@ -591,7 +653,7 @@ void run(Header &data,
   header.datatype() = DataType::from<T>();
   auto output = Image<T>::create(output_name, header);
   // run
-  DenoisingFunctor<T> func(data.size(3), kernel, mask, noise, rank, max_dist, voxels, estimator);
+  DenoisingFunctor<T> func(data.size(3), kernel, filter, mask, noise, rank, sum_weights, max_dist, voxels, estimator);
   ThreadedLoop("running MP-PCA denoising", data, 0, 3).run(func, input, output);
 }
 
@@ -613,6 +675,11 @@ void run() {
   if (!opt.empty())
     estimator = estimator_type(int(opt[0][0]));
 
+  filter_type filter = filter_type::FROBENIUS;
+  opt = get_options("filter");
+  if (!opt.empty())
+    filter = filter_type(int(opt[0][0]));
+
   Image<real_type> noise;
   opt = get_options("noise");
   if (!opt.empty()) {
@@ -632,6 +699,21 @@ void run() {
     rank = Image<uint16_t>::create(opt[0][0], header);
   }
 
+  Image<float> sum_weights;
+  opt = get_options("sumweights");
+  if (!opt.empty()) {
+    Header header(dwi);
+    header.ndim() = 3;
+    header.datatype() = DataType::Float32;
+    header.datatype().set_byte_order_native();
+    header.reset_intensity_scaling();
+    sum_weights = Image<float>::create(opt[0][0], header);
+    if (filter == filter_type::TRUNCATE) {
+      WARN("Note that with a truncation filter, "
+           "output image from -sumweights option will be equivalent to rank");
+    }
+  }
+
   Image<float> max_dist;
   opt = get_options("max_dist");
   if (!opt.empty()) {
@@ -714,19 +796,19 @@ void run() {
   switch (prec) {
   case 0:
     INFO("select real float32 for processing");
-    run<float>(dwi, mask, noise, rank, max_dist, voxels, argument[1], kernel, estimator);
+    run<float>(dwi, mask, noise, rank, sum_weights, max_dist, voxels, argument[1], kernel, filter, estimator);
     break;
   case 1:
     INFO("select real float64 for processing");
-    run<double>(dwi, mask, noise, rank, max_dist, voxels, argument[1], kernel, estimator);
+    run<double>(dwi, mask, noise, rank, sum_weights, max_dist, voxels, argument[1], kernel, filter, estimator);
     break;
   case 2:
     INFO("select complex float32 for processing");
-    run<cfloat>(dwi, mask, noise, rank, max_dist, voxels, argument[1], kernel, estimator);
+    run<cfloat>(dwi, mask, noise, rank, sum_weights, max_dist, voxels, argument[1], kernel, filter, estimator);
     break;
   case 3:
     INFO("select complex float64 for processing");
-    run<cdouble>(dwi, mask, noise, rank, max_dist, voxels, argument[1], kernel, estimator);
+    run<cdouble>(dwi, mask, noise, rank, sum_weights, max_dist, voxels, argument[1], kernel, filter, estimator);
     break;
   }
 }

From aee5c06a7437328a22f9785e5b4f1c889ed4ec79 Mon Sep 17 00:00:00 2001
From: Robert Smith <robert.smith@florey.edu.au>
Date: Mon, 11 Nov 2024 22:26:00 +1100
Subject: [PATCH 14/34] dwidenoise: Add new estimator

Closes #2591.
---
 cmd/dwidenoise.cpp | 180 ++++++++++++++++++++++++++++++++-------------
 1 file changed, 128 insertions(+), 52 deletions(-)

diff --git a/cmd/dwidenoise.cpp b/cmd/dwidenoise.cpp
index 493fae334f..53fd29cc04 100644
--- a/cmd/dwidenoise.cpp
+++ b/cmd/dwidenoise.cpp
@@ -24,8 +24,8 @@ using namespace MR;
 using namespace App;
 
 const std::vector<std::string> dtypes = {"float32", "float64"};
-const std::vector<std::string> estimators = {"exp1", "exp2"};
-enum class estimator_type { EXP1, EXP2 };
+const std::vector<std::string> estimators = {"exp1", "exp2", "mrm2022"};
+enum class estimator_type { EXP1, EXP2, MRM2022 };
 
 const std::vector<std::string> shapes = {"cuboid", "sphere"};
 enum class shape_type { CUBOID, SPHERE };
@@ -45,10 +45,8 @@ void usage() {
     " using the prior knowledge that the eigenspectrum of random covariance matrices"
     " is described by the universal Marchenko-Pastur (MP) distribution."
     " Fitting the MP distribution to the spectrum of patch-wise signal matrices"
-    " hence provides an estimator of the noise level 'sigma',"
-    " as was first shown in Veraart et al. (2016)"
-    " and later improved in Cordero-Grande et al. (2019)."
-    " This noise level estimate then determines the optimal cut-off for PCA denoising."
+    " hence provides an estimator of the noise level 'sigma';"
+    " this noise level estimate then determines the optimal cut-off for PCA denoising."
 
   + "Important note:"
     " image denoising must be performed as the first step of the image processing pipeline."
@@ -103,7 +101,12 @@ void usage() {
 
   + "Cordero-Grande, L.; Christiaens, D.; Hutter, J.; Price, A.N.; Hajnal, J.V. " // Internal
     "Complex diffusion-weighted image estimation via matrix recovery under general noise models. "
-    "NeuroImage, 2019, 200, 391-404, doi: 10.1016/j.neuroimage.2019.06.039";
+    "NeuroImage, 2019, 200, 391-404, doi: 10.1016/j.neuroimage.2019.06.039"
+
+  + "* If using -estimator mrm2022: "
+    "Olesen, J.L.; Ianus, A.; Ostergaard, L.; Shemesh, N.; Jespersen, S.N. "
+    "Tensor denoising of multidimensional MRI data. "
+    "Magnetic Resonance in Medicine, 2022, 89(3), 1160-1172";
 
   ARGUMENTS
   + Argument("dwi", "the input diffusion-weighted image.").type_image_in()
@@ -124,9 +127,10 @@ void usage() {
            "Select the noise level estimator"
            " (default = Exp2),"
            " either: \n"
-           "* Exp1: the original estimator used in Veraart et al. (2016), or \n"
-           "* Exp2: the improved estimator introduced in Cordero-Grande et al. (2019).")
-    + Argument("Exp1/Exp2").type_choice(estimators)
+           "* Exp1: the original estimator used in Veraart et al. (2016); \n"
+           "* Exp2: the improved estimator introduced in Cordero-Grande et al. (2019); \n"
+           "* MRM2022: the alternative estimator introduced in Olesen et al. (2022).")
+    + Argument("algorithm").type_choice(estimators)
   + Option("filter",
            "Modulate how components are filtered based on their eigenvalues; "
            "options are: " + join(filters, ",") + "; default: frobenius")
@@ -430,6 +434,89 @@ class KernelSphereFixedRadius : public KernelSphereBase {
   const ssize_t maximum_size;
 };
 
+class EstimatorResult {
+public:
+  EstimatorResult() : cutoff_p(0), sigma2(0.0) {}
+  ssize_t cutoff_p;
+  double sigma2;
+};
+
+class EstimatorBase {
+public:
+  EstimatorBase() = default;
+  virtual EstimatorResult operator()(const vector_type &eigenvalues, const ssize_t m, const ssize_t n) const = 0;
+};
+
+template <ssize_t version> class EstimatorExp : public EstimatorBase {
+public:
+  EstimatorExp() = default;
+  EstimatorResult operator()(const vector_type &s, const ssize_t m, const ssize_t n) const final {
+    EstimatorResult result;
+    const ssize_t r = std::min(m, n);
+    const ssize_t q = std::max(m, n);
+    assert(s.size() == r);
+    const double lam_r = std::max(s[0], 0.0) / q;
+    double clam = 0.0;
+    for (ssize_t p = 0; p < r; ++p) // p+1 is the number of noise components
+    {                               // (as opposed to the paper where p is defined as the number of signal components)
+      const double lam = std::max(s[p], 0.0) / q;
+      clam += lam;
+      double denominator = std::numeric_limits<double>::signaling_NaN();
+      switch (version) {
+      case 1:
+        denominator = q;
+        break;
+      case 2:
+        denominator = q - (r - p - 1);
+        break;
+      default:
+        assert(false);
+      }
+      const double gam = double(p + 1) / denominator;
+      const double sigsq1 = clam / double(p + 1);
+      const double sigsq2 = (lam - lam_r) / (4.0 * std::sqrt(gam));
+      // sigsq2 > sigsq1 if signal else noise
+      if (sigsq2 < sigsq1) {
+        result.sigma2 = sigsq1;
+        result.cutoff_p = p + 1;
+      }
+    }
+    return result;
+  }
+};
+
+class EstimatorMRM2022 : public EstimatorBase {
+public:
+  EstimatorMRM2022() = default;
+  EstimatorResult operator()(const vector_type &s, const ssize_t m, const ssize_t n) const final {
+    EstimatorResult result;
+    const ssize_t mprime = std::min(m, n);
+    const ssize_t nprime = std::max(m, n);
+    const double sigmasq_to_lamplus = Math::pow2(std::sqrt(nprime) + std::sqrt(mprime));
+    assert(s.size() == mprime);
+    double clam = 0.0;
+    for (ssize_t i = 0; i != mprime; ++i)
+      clam += std::max(s[i], 0.0);
+    clam /= nprime;
+    // Unlike Exp# code,
+    //   MRM2022 article uses p to index number of signal components,
+    //   and here doing a direct translation of the manuscript content to code
+    double lamplusprev = -std::numeric_limits<double>::infinity();
+    for (ssize_t p = 0; p < mprime; ++p) {
+      const ssize_t i = mprime - 1 - p;
+      const double lam = std::max(s[i], 0.0) / nprime;
+      if (lam < lamplusprev)
+        return result;
+      clam -= lam;
+      const double sigmasq = clam / ((mprime - p) * (nprime - p));
+      lamplusprev = sigmasq * sigmasq_to_lamplus;
+      result.cutoff_p = i;
+      result.sigma2 = sigmasq;
+    }
+    return result;
+  }
+};
+
 template <typename F> class DenoisingFunctor {
 
 public:
@@ -444,7 +531,7 @@ template <typename F> class DenoisingFunctor {
                    Image<float> &sum_weights,
                    Image<float> &max_dist,
                    Image<uint16_t> &voxels,
-                   estimator_type estimator)
+                   std::shared_ptr<EstimatorBase> estimator)
       : kernel(kernel),
         filter(filter),
         m(ndwi),
@@ -512,49 +599,26 @@ template <typename F> class DenoisingFunctor {
     // eigenvalues sorted in increasing order:
     s.head(r) = eig.eigenvalues().template cast<double>();
 
-    // Marchenko-Pastur optimal threshold
-    const double lam_r = std::max(s[0], 0.0) / q;
-    double sigma2 = 0.0;
-    ssize_t cutoff_p = 0;
-    for (ssize_t p = 0; p < r; ++p) // p+1 is the number of noise components
-    {                               // (as opposed to the paper where p is defined as the number of signal components)
-      const double lam = std::max(s[p], 0.0) / q;
-      clam[p] = (p == 0 ? 0.0 : clam[p - 1]) + lam;
-      double denominator = std::numeric_limits<double>::signaling_NaN();
-      switch (estimator) {
-      case estimator_type::EXP1:
-        denominator = q;
-        break;
-      case estimator_type::EXP2:
-        denominator = q - (r - p - 1);
-        break;
-      default:
-        assert(false);
-      }
-      const double gam = double(p + 1) / denominator;
-      const double sigsq1 = clam[p] / double(p + 1);
-      const double sigsq2 = (lam - lam_r) / (4.0 * std::sqrt(gam));
-      // sigsq2 > sigsq1 if signal else noise
-      if (sigsq2 < sigsq1) {
-        sigma2 = sigsq1;
-        cutoff_p = p + 1;
-      }
-    }
+    // Marchenko-Pastur optimal threshold determination
+    const EstimatorResult threshold = (*estimator)(s, m, n);
 
     // Generate weights vector
     double sum_weights = 0.0;
     switch (filter) {
     case filter_type::TRUNCATE:
-      w.head(cutoff_p).setZero();
-      w.segment(cutoff_p, r - cutoff_p).setOnes();
-      sum_weights = r - cutoff_p;
+      w.head(threshold.cutoff_p).setZero();
+      w.segment(threshold.cutoff_p, r - threshold.cutoff_p).setOnes();
+      sum_weights = r - threshold.cutoff_p;
       break;
     case filter_type::FROBENIUS: {
       const double beta = r / q;
-      const double threshold = 1.0 + std::sqrt(beta);
+      const double transition = 1.0 + std::sqrt(beta);
+      double clam = 0.0;
       for (ssize_t i = 0; i != r; ++i) {
-        const double y = clam[i] / (sigma2 * (i + 1));
-        const double nu = y > threshold ? std::sqrt(Math::pow2(Math::pow2(y) - beta - 1.0) - (4.0 * beta)) / y : 0.0;
+        const double lam = std::max(s[i], 0.0) / q;
+        clam += lam;
+        const double y = clam / (threshold.sigma2 * (i + 1));
+        const double nu = y > transition ? std::sqrt(Math::pow2(Math::pow2(y) - beta - 1.0) - (4.0 * beta)) / y : 0.0;
         w[i] = nu / y;
         sum_weights += w[i];
       }
@@ -580,11 +644,11 @@ template <typename F> class DenoisingFunctor {
     // Store additional output maps if requested
     if (noise.valid()) {
       assign_pos_of(dwi, 0, 3).to(noise);
-      noise.value() = real_type(std::sqrt(sigma2));
+      noise.value() = real_type(std::sqrt(threshold.sigma2));
     }
     if (rankmap.valid()) {
       assign_pos_of(dwi, 0, 3).to(rankmap);
-      rankmap.value() = uint16_t(r - cutoff_p);
+      rankmap.value() = uint16_t(r - threshold.cutoff_p);
     }
     if (sumweightsmap.valid()) {
       assign_pos_of(dwi, 0, 3).to(sumweightsmap);
@@ -605,7 +669,7 @@ template <typename F> class DenoisingFunctor {
   std::shared_ptr<KernelBase> kernel;
   filter_type filter;
   const ssize_t m;
-  const estimator_type estimator;
+  std::shared_ptr<EstimatorBase> estimator;
 
   // Reusable memory
   MatrixType X;
@@ -646,7 +710,7 @@ void run(Header &data,
          const std::string &output_name,
          std::shared_ptr<KernelBase> kernel,
          filter_type filter,
-         estimator_type estimator) {
+         std::shared_ptr<EstimatorBase> estimator) {
   auto input = data.get_image<T>().with_direct_io(3);
   // create output
   Header header(data);
@@ -670,10 +734,22 @@ void run() {
     check_dimensions(mask, dwi, 0, 3);
   }
 
-  estimator_type estimator = estimator_type::EXP2; // default: Exp2 (unbiased estimator)
+  std::shared_ptr<EstimatorBase> estimator;
   opt = get_options("estimator");
-  if (!opt.empty())
-    estimator = estimator_type(int(opt[0][0]));
+  const estimator_type est = opt.empty() ? estimator_type::EXP2 : estimator_type((int)(opt[0][0]));
+  switch (est) {
+  case estimator_type::EXP1:
+    estimator = std::make_shared<EstimatorExp<1>>();
+    break;
+  case estimator_type::EXP2:
+    estimator = std::make_shared<EstimatorExp<2>>();
+    break;
+  case estimator_type::MRM2022:
+    estimator = std::make_shared<EstimatorMRM2022>();
+    break;
+  default:
+    assert(false);
+  }
 
   filter_type filter = filter_type::FROBENIUS;
   opt = get_options("filter");

From f59b78d17ef99a063854e0d7cf19acd6da67f132 Mon Sep 17 00:00:00 2001
From: Robert Smith <robert.smith@florey.edu.au>
Date: Wed, 13 Nov 2024 13:53:18 +1100
Subject: [PATCH 15/34] dwidenoise: Add overcomplete local PCA

Default behaviour is to use a Gaussian kernel (as used in Cordero-Grande et al. 2019) with FWHM equal to twice the voxel spacing.
Closes #3024.
---
 cmd/dwidenoise.cpp | 263 ++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 227 insertions(+), 36 deletions(-)

diff --git a/cmd/dwidenoise.cpp b/cmd/dwidenoise.cpp
index 53fd29cc04..3f29425dec 100644
--- a/cmd/dwidenoise.cpp
+++ b/cmd/dwidenoise.cpp
@@ -14,6 +14,8 @@
  * For more details, see http://www.mrtrix.org/.
  */
 
+#include <mutex>
+
 #include "command.h"
 #include "image.h"
 
@@ -34,6 +36,9 @@ constexpr default_type sphere_multiplier_default = 1.0 / 0.85;
 const std::vector<std::string> filters = {"truncate", "frobenius"};
 enum class filter_type { TRUNCATE, FROBENIUS };
 
+const std::vector<std::string> aggregators = {"exclusive", "Gaussian", "invL0", "rank", "uniform"};
+enum class aggregator_type { EXCLUSIVE, GAUSSIAN, INVL0, RANK, UNIFORM };
+
 // clang-format off
 void usage() {
 
@@ -83,7 +88,25 @@ void usage() {
     "will be used to attenuate eigenvectors based on the estimated noise level. "
     "Hard truncation of sub-threshold components"
     "---which was the behaviour of the dwidenoise command in version 3.0.x---"
-    "can be activated using -filter truncate.";
+    "can be activated using -filter truncate."
+
+  + "-aggregation exclusive corresponds to the behaviour of the dwidenoise command in version 3.0.x, "
+    "where the output intensities for a given image voxel are determined exclusively "
+    "from the PCA decomposition where the sliding spatial window is centred at that voxel. "
+    "In all other use cases, so-called \"overcomplete local PCA\" is performed, "
+    "where the intensities for an output image voxel are some combination of all PCA decompositions "
+    "for which that voxel is included in the local spatial kernel. "
+    "There are multiple algebraic forms that modulate the weight with which each decomposition "
+    "contributes with greater or lesser strength toward the output image intensities. "
+    "The various options are: "
+    "'Gaussian': A Gaussian distribution with FWHM equal to twice the voxel size, "
+      "such that decompisitions centred more closely to the output voxel have greater influence; "
+    "'invL0': The inverse of the L0 norm (ie. rank) of each decomposition, "
+      "as used in Manjon et al. 2013; "
+    "'rank': The rank of each decomposition, "
+      "such that high-rank decompositions contribute more strongly to the output intensities "
+      "regardless of distance between the output voxel and the centre of the decomposition kernel; "
+    "'uniform': All decompositions that include the output voxel in the sliding spatial window contribute equally.";
 
   AUTHOR = "Daan Christiaens (daan.christiaens@kcl.ac.uk)"
            " and Jelle Veraart (jelle.veraart@nyumc.org)"
@@ -106,7 +129,12 @@ void usage() {
   + "* If using -estimator mrm2022: "
     "Olesen, J.L.; Ianus, A.; Ostergaard, L.; Shemesh, N.; Jespersen, S.N. "
     "Tensor denoising of multidimensional MRI data. "
-    "Magnetic Resonance in Medicine, 2022, 89(3), 1160-1172";
+    "Magnetic Resonance in Medicine, 2022, 89(3), 1160-1172"
+
+  + "* If using anything other than -aggregation exclusive: "
+    "Manjon, J.V.; Coupe, P.; Concha, L.; Buades, A.; D. Collins, D.L.; Robles, M. "
+    "Diffusion Weighted Image Denoising Using Overcomplete Local PCA. "
+    "PLoS ONE, 2013, 8(9), e73021";
 
   ARGUMENTS
   + Argument("dwi", "the input diffusion-weighted image.").type_image_in()
@@ -135,6 +163,13 @@ void usage() {
            "Modulate how components are filtered based on their eigenvalues; "
            "options are: " + join(filters, ",") + "; default: frobenius")
     + Argument("choice").type_choice(filters)
+  + Option("aggregator",
+           "Select how the outcomes of multiple PCA outcomes centred at different voxels "
+           "contribute to the reconstructed DWI signal in each voxel; "
+           "options are: " + join(aggregators, ",") + "; default: Gaussian")
+    + Argument("choice").type_choice(aggregators)
+  // TODO For specifically the Gaussian aggregator,
+  //   should ideally be possible to select the FWHM of the aggregator
 
   + OptionGroup("Options for exporting additional data regarding PCA behaviour")
   + Option("noise",
@@ -156,6 +191,9 @@ void usage() {
   + Option("voxels",
            "The number of voxels that contributed to the PCA for processing of each voxel")
     + Argument("image").type_image_out()
+  + Option("aggregation_sum",
+           "The sum of aggregation weights of those patches contributing to each output voxel")
+    + Argument("image").type_image_out()
 
   + OptionGroup("Options for controlling the sliding spatial window")
   + Option("shape",
@@ -204,7 +242,6 @@ void usage() {
 }
 // clang-format on
 
-using real_type = float;
 using voxel_type = Eigen::Array<int, 3, 1>;
 using vector_type = Eigen::VectorXd;
 
@@ -226,6 +263,9 @@ class KernelVoxel {
   }
   bool operator<(const KernelVoxel &that) const { return sq_distance < that.sq_distance; }
   default_type distance() const { return std::sqrt(sq_distance); }
+  // TODO Sometimes this acts as an offset, other times it acts as an absolute voxel index
+  // Consider either renaming, or actually using two different classes
+  // The latter could use ssize_t instead of int to better indicate this
   voxel_type offset;
   default_type sq_distance;
 };
@@ -454,7 +494,6 @@ template <ssize_t version> class EstimatorExp : public EstimatorBase {
     EstimatorResult result;
     const ssize_t r = std::min(m, n);
     const ssize_t q = std::max(m, n);
-    assert(s.size() == r);
     const double lam_r = std::max(s[0], 0.0) / q;
     double clam = 0.0;
     for (ssize_t p = 0; p < r; ++p) // p+1 is the number of noise components
@@ -493,7 +532,6 @@ class EstimatorMRM2022 : public EstimatorBase {
     const ssize_t mprime = std::min(m, n);
     const ssize_t nprime = std::max(m, n);
     const double sigmasq_to_lamplus = Math::pow2(std::sqrt(nprime) + std::sqrt(mprime));
-    assert(s.size() == mprime);
     double clam = 0.0;
     for (ssize_t i = 0; i != mprime; ++i)
       clam += std::max(s[i], 0.0);
@@ -522,32 +560,40 @@ template <typename F> class DenoisingFunctor {
 public:
   using MatrixType = Eigen::Matrix<F, Eigen::Dynamic, Eigen::Dynamic>;
 
-  DenoisingFunctor(int ndwi,
+  DenoisingFunctor(const Header &header,
                    std::shared_ptr<KernelBase> kernel,
                    filter_type filter,
+                   aggregator_type aggregator,
                    Image<bool> &mask,
-                   Image<real_type> &noise,
+                   Image<float> &noise,
                    Image<uint16_t> &rank,
                    Image<float> &sum_weights,
                    Image<float> &max_dist,
                    Image<uint16_t> &voxels,
+                   // TODO Would be preferable for this to be double if computations are happening using double
+                   Image<float> &aggregation_weight_map,
                    std::shared_ptr<EstimatorBase> estimator)
       : kernel(kernel),
         filter(filter),
-        m(ndwi),
+        aggregator(aggregator),
+        // FWHM = 2 x cube root of voxel spacings
+        gaussian_multiplier(-std::log(2.0) /
+                            Math::pow2(std::cbrt(header.spacing(0) * header.spacing(1) * header.spacing(2)))),
+        m(header.size(3)),
         estimator(estimator),
-        X(ndwi, kernel->estimated_size()),
+        mask(mask),
+        X(m, kernel->estimated_size()),
         XtX(std::min(m, kernel->estimated_size()), std::min(m, kernel->estimated_size())),
         eig(std::min(m, kernel->estimated_size())),
         s(std::min(m, kernel->estimated_size())),
         clam(std::min(m, kernel->estimated_size())),
         w(std::min(m, kernel->estimated_size())),
-        mask(mask),
         noise(noise),
         rankmap(rank),
         sumweightsmap(sum_weights),
         maxdistmap(max_dist),
-        voxelsmap(voxels) {}
+        voxelsmap(voxels),
+        aggregation_weight_map(aggregation_weight_map) {}
 
   template <typename ImageType> void operator()(ImageType &dwi, ImageType &out) {
     // Process voxels in mask only
@@ -627,24 +673,68 @@ template <typename F> class DenoisingFunctor {
       assert(false);
     }
 
-    // recombine data using only eigenvectors above threshold:
-    if (m <= n)
-      X.col(neighbourhood.centre_index) =
-          eig.eigenvectors() *
-          (w.head(r).cast<F>().asDiagonal() * (eig.eigenvectors().adjoint() * X.col(neighbourhood.centre_index)));
-    else
-      X.col(neighbourhood.centre_index) =
-          X.leftCols(n) * (eig.eigenvectors() * (w.head(r).cast<F>().asDiagonal() *
-                                                 eig.eigenvectors().adjoint().col(neighbourhood.centre_index)));
-
-    // Store output
-    assign_pos_of(dwi).to(out);
-    out.row(3) = X.col(neighbourhood.centre_index);
+    // recombine data using only eigenvectors above threshold
+    // If only the data computed when this voxel was the centre of the patch
+    //   is to be used for synthesis of the output image,
+    //   then only that individual column needs to be reconstructed;
+    //   if however the result from this patch is to contribute to the synthesized image
+    //   for all voxels that were utilised within this patch,
+    //   then we need to instead compute the full projection
+    switch (aggregator) {
+    case aggregator_type::EXCLUSIVE:
+      if (m <= n)
+        X.col(neighbourhood.centre_index) =
+            eig.eigenvectors() *
+            (w.head(r).cast<F>().asDiagonal() * (eig.eigenvectors().adjoint() * X.col(neighbourhood.centre_index)));
+      else
+        X.col(neighbourhood.centre_index) =
+            X.leftCols(n) * (eig.eigenvectors() * (w.head(r).cast<F>().asDiagonal() *
+                                                   eig.eigenvectors().adjoint().col(neighbourhood.centre_index)));
+      assign_pos_of(dwi).to(out);
+      out.row(3) = X.col(neighbourhood.centre_index);
+      if (aggregation_weight_map.valid()) {
+        assign_pos_of(dwi, 0, 3).to(aggregation_weight_map);
+        aggregation_weight_map.value() = 1.0;
+      }
+      break;
+    default: {
+      if (m <= n)
+        X = eig.eigenvectors() * (w.head(r).cast<F>().asDiagonal() * (eig.eigenvectors().adjoint() * X));
+      else
+        X.leftCols(n) =
+            X.leftCols(n) * (eig.eigenvectors() * (w.head(r).cast<F>().asDiagonal() * eig.eigenvectors().adjoint()));
+      std::lock_guard<std::mutex> lock(mutex_aggregator);
+      for (size_t voxel_index = 0; voxel_index != neighbourhood.voxels.size(); ++voxel_index) {
+        assign_pos_of(neighbourhood.voxels[voxel_index].offset, 0, 3).to(out);
+        assign_pos_of(neighbourhood.voxels[voxel_index].offset).to(aggregation_weight_map);
+        double weight = std::numeric_limits<double>::signaling_NaN();
+        switch (aggregator) {
+        case aggregator_type::EXCLUSIVE:
+          assert(false);
+          break;
+        case aggregator_type::GAUSSIAN:
+          weight = std::exp(gaussian_multiplier * neighbourhood.voxels[voxel_index].sq_distance);
+          break;
+        case aggregator_type::INVL0:
+          weight = 1.0 / (1 + r - threshold.cutoff_p);
+          break;
+        case aggregator_type::RANK:
+          weight = r - threshold.cutoff_p;
+          break;
+        case aggregator_type::UNIFORM:
+          weight = 1.0;
+          break;
+        }
+        out.row(3) += weight * X.col(voxel_index);
+        aggregation_weight_map.value() += weight;
+      }
+    } break;
+    }
 
     // Store additional output maps if requested
     if (noise.valid()) {
       assign_pos_of(dwi, 0, 3).to(noise);
-      noise.value() = real_type(std::sqrt(threshold.sigma2));
+      noise.value() = float(std::sqrt(threshold.sigma2));
     }
     if (rankmap.valid()) {
       assign_pos_of(dwi, 0, 3).to(rankmap);
@@ -656,7 +746,7 @@ template <typename F> class DenoisingFunctor {
     }
     if (maxdistmap.valid()) {
       assign_pos_of(dwi, 0, 3).to(maxdistmap);
-      maxdistmap.value() = float(neighbourhood.max_distance);
+      maxdistmap.value() = neighbourhood.max_distance;
     }
     if (voxelsmap.valid()) {
       assign_pos_of(dwi, 0, 3).to(voxelsmap);
@@ -668,8 +758,11 @@ template <typename F> class DenoisingFunctor {
   // Denoising configuration
   std::shared_ptr<KernelBase> kernel;
   filter_type filter;
+  aggregator_type aggregator;
+  double gaussian_multiplier;
   const ssize_t m;
   std::shared_ptr<EstimatorBase> estimator;
+  Image<bool> mask;
 
   // Reusable memory
   MatrixType X;
@@ -679,15 +772,22 @@ template <typename F> class DenoisingFunctor {
   vector_type clam;
   vector_type w;
 
-  Image<bool> mask;
+  // Data that can only be written in a thread-safe manner
+  // Note that this applies not just to this scratch buffer, but also the output image
+  //   (while it would be thread-safe to create a full copy of the output image for each thread
+  //   and combine them only at destruction time,
+  //   this runs the risk of becoming prohibitively large)
+  // Not placing this within a MutexProtexted<> as the image type is still templated
+  static std::mutex mutex_aggregator;
 
   // Export images
   // TODO Group these into a class?
-  Image<real_type> noise;
+  Image<float> noise;
   Image<uint16_t> rankmap;
   Image<float> sumweightsmap;
   Image<float> maxdistmap;
   Image<uint16_t> voxelsmap;
+  Image<float> aggregation_weight_map;
 
   template <typename ImageType> void load_data(ImageType &image, const std::vector<KernelVoxel> &voxels) {
     const voxel_type pos({int(image.index(0)), int(image.index(1)), int(image.index(2))});
@@ -698,18 +798,26 @@ template <typename F> class DenoisingFunctor {
     assign_pos_of(pos, 0, 3).to(image);
   }
 };
+template <typename F> std::mutex DenoisingFunctor<F>::mutex_aggregator;
+
+// Necessary to allow normalisation by sum of aggregation weights
+//   where the image type is cdouble, but aggregation weights are float
+// (operations combining complex & real types not allowed to be of different precision)
+std::complex<double> operator/(const std::complex<double> &c, const float n) { return c / double(n); }
 
 template <typename T>
 void run(Header &data,
          Image<bool> &mask,
-         Image<real_type> &noise,
+         Image<float> &noise,
          Image<uint16_t> &rank,
          Image<float> &sum_weights,
          Image<float> &max_dist,
          Image<uint16_t> &voxels,
+         Image<float> &aggregation_sum,
          const std::string &output_name,
          std::shared_ptr<KernelBase> kernel,
          filter_type filter,
+         aggregator_type aggregator,
          std::shared_ptr<EstimatorBase> estimator) {
   auto input = data.get_image<T>().with_direct_io(3);
   // create output
@@ -717,8 +825,16 @@ void run(Header &data,
   header.datatype() = DataType::from<T>();
   auto output = Image<T>::create(output_name, header);
   // run
-  DenoisingFunctor<T> func(data.size(3), kernel, filter, mask, noise, rank, sum_weights, max_dist, voxels, estimator);
+  DenoisingFunctor<T> func(
+      data, kernel, filter, aggregator, mask, noise, rank, sum_weights, max_dist, voxels, aggregation_sum, estimator);
   ThreadedLoop("running MP-PCA denoising", data, 0, 3).run(func, input, output);
+  // Rescale output if performing aggregation
+  if (aggregator == aggregator_type::EXCLUSIVE)
+    return;
+  for (auto l_voxel = Loop(aggregation_sum)(output, aggregation_sum); l_voxel; ++l_voxel) {
+    for (auto l_volume = Loop(3)(output); l_volume; ++l_volume)
+      output.value() /= float(aggregation_sum.value());
+  }
 }
 
 void run() {
@@ -756,13 +872,18 @@ void run() {
   if (!opt.empty())
     filter = filter_type(int(opt[0][0]));
 
-  Image<real_type> noise;
+  aggregator_type aggregator = aggregator_type::GAUSSIAN;
+  opt = get_options("aggregator");
+  if (!opt.empty())
+    aggregator = aggregator_type(int(opt[0][0]));
+
+  Image<float> noise;
   opt = get_options("noise");
   if (!opt.empty()) {
     Header header(dwi);
     header.ndim() = 3;
     header.datatype() = DataType::Float32;
-    noise = Image<real_type>::create(opt[0][0], header);
+    noise = Image<float>::create(opt[0][0], header);
   }
 
   Image<uint16_t> rank;
@@ -811,6 +932,28 @@ void run() {
     voxels = Image<uint16_t>::create(opt[0][0], header);
   }
 
+  Image<float> aggregation_sum;
+  Header header_aggregation(dwi);
+  header_aggregation.ndim() = 3;
+  header_aggregation.datatype() = DataType::Float64;
+  header_aggregation.datatype().set_byte_order_native();
+  header_aggregation.reset_intensity_scaling();
+  opt = get_options("aggregation_sum");
+  if (!opt.empty()) {
+    if (aggregator == aggregator_type::EXCLUSIVE) {
+      WARN("Output from -aggregation_sum will just contain 1 for every voxel processed: "
+           "no patch aggregation takes place when output series comex exclusively from central patch");
+    }
+    Header header(dwi);
+    header.ndim() = 3;
+    header.datatype() = DataType::Float32;
+    header.datatype().set_byte_order_native();
+    header.reset_intensity_scaling();
+    aggregation_sum = Image<float>::create(opt[0][0], header_aggregation);
+  } else if (aggregator != aggregator_type::EXCLUSIVE) {
+    aggregation_sum = Image<float>::scratch(header_aggregation, "Scratch buffer for patch aggregation weights");
+  }
+
   opt = get_options("shape");
   const shape_type shape = opt.empty() ? shape_type::SPHERE : shape_type((int)(opt[0][0]));
   std::shared_ptr<KernelBase> kernel;
@@ -872,19 +1015,67 @@ void run() {
   switch (prec) {
   case 0:
     INFO("select real float32 for processing");
-    run<float>(dwi, mask, noise, rank, sum_weights, max_dist, voxels, argument[1], kernel, filter, estimator);
+    run<float>(dwi,
+               mask,
+               noise,
+               rank,
+               sum_weights,
+               max_dist,
+               voxels,
+               aggregation_sum,
+               argument[1],
+               kernel,
+               filter,
+               aggregator,
+               estimator);
     break;
   case 1:
     INFO("select real float64 for processing");
-    run<double>(dwi, mask, noise, rank, sum_weights, max_dist, voxels, argument[1], kernel, filter, estimator);
+    run<double>(dwi,
+                mask,
+                noise,
+                rank,
+                sum_weights,
+                max_dist,
+                voxels,
+                aggregation_sum,
+                argument[1],
+                kernel,
+                filter,
+                aggregator,
+                estimator);
     break;
   case 2:
     INFO("select complex float32 for processing");
-    run<cfloat>(dwi, mask, noise, rank, sum_weights, max_dist, voxels, argument[1], kernel, filter, estimator);
+    run<cfloat>(dwi,
+                mask,
+                noise,
+                rank,
+                sum_weights,
+                max_dist,
+                voxels,
+                aggregation_sum,
+                argument[1],
+                kernel,
+                filter,
+                aggregator,
+                estimator);
     break;
   case 3:
     INFO("select complex float64 for processing");
-    run<cdouble>(dwi, mask, noise, rank, sum_weights, max_dist, voxels, argument[1], kernel, filter, estimator);
+    run<cdouble>(dwi,
+                 mask,
+                 noise,
+                 rank,
+                 sum_weights,
+                 max_dist,
+                 voxels,
+                 aggregation_sum,
+                 argument[1],
+                 kernel,
+                 filter,
+                 aggregator,
+                 estimator);
     break;
   }
 }

From 0e015a25228e2cac526f533f2dbb7dd054570dfe Mon Sep 17 00:00:00 2001
From: Robert Smith <robert.smith@florey.edu.au>
Date: Wed, 13 Nov 2024 16:11:21 +1100
Subject: [PATCH 16/34] dwidenoise: New option -weightedrank

An estimated PCA rank that is the weighted average of the patches that contributed to the rank estimation via overcomplete local PCA reconstruction is likely a superior measure of signal rank for downstream use.
---
 cmd/dwidenoise.cpp | 100 ++++++++++++++++++++++++++++++++-------------
 1 file changed, 71 insertions(+), 29 deletions(-)

diff --git a/cmd/dwidenoise.cpp b/cmd/dwidenoise.cpp
index 3f29425dec..a99bfd8806 100644
--- a/cmd/dwidenoise.cpp
+++ b/cmd/dwidenoise.cpp
@@ -36,7 +36,7 @@ constexpr default_type sphere_multiplier_default = 1.0 / 0.85;
 const std::vector<std::string> filters = {"truncate", "frobenius"};
 enum class filter_type { TRUNCATE, FROBENIUS };
 
-const std::vector<std::string> aggregators = {"exclusive", "Gaussian", "invL0", "rank", "uniform"};
+const std::vector<std::string> aggregators = {"exclusive", "gaussian", "invl0", "rank", "uniform"};
 enum class aggregator_type { EXCLUSIVE, GAUSSIAN, INVL0, RANK, UNIFORM };
 
 // clang-format off
@@ -171,6 +171,10 @@ void usage() {
   // TODO For specifically the Gaussian aggregator,
   //   should ideally be possible to select the FWHM of the aggregator
 
+  // TODO Consider renaming some options to better distinguish between:
+  // - Parameters arising from PCA-based noise level estimation
+  // - Parameters encoding properties of the output data
+
   + OptionGroup("Options for exporting additional data regarding PCA behaviour")
   + Option("noise",
            "The output noise map,"
@@ -180,10 +184,13 @@ void usage() {
            " so a scale factor sqrt(2) applies.")
     + Argument("image").type_image_out()
   + Option("rank",
-           "The selected signal rank of the output denoised image.")
+           "The estimated signal rank for the denoising patch centred at each voxel")
+    + Argument("image").type_image_out()
+  + Option("weightedrank",
+           "The weighted mean rank for the output image data, accounting for multi-patch aggregation")
     + Argument("image").type_image_out()
   + Option("sumweights",
-           "the sum of eigenvector weights contributed to the output image")
+           "the sum of eigenvector weights computed for the denoising patch centred at each voxel")
     + Argument("image").type_image_out()
   + Option("max_dist",
            "The maximum distance between a voxel and another voxel that was included in the local denoising patch")
@@ -567,6 +574,7 @@ template <typename F> class DenoisingFunctor {
                    Image<bool> &mask,
                    Image<float> &noise,
                    Image<uint16_t> &rank,
+                   Image<float> &weighted_rank,
                    Image<float> &sum_weights,
                    Image<float> &max_dist,
                    Image<uint16_t> &voxels,
@@ -590,6 +598,7 @@ template <typename F> class DenoisingFunctor {
         w(std::min(m, kernel->estimated_size())),
         noise(noise),
         rankmap(rank),
+        weightedrankmap(weighted_rank),
         sumweightsmap(sum_weights),
         maxdistmap(max_dist),
         voxelsmap(voxels),
@@ -696,6 +705,10 @@ template <typename F> class DenoisingFunctor {
         assign_pos_of(dwi, 0, 3).to(aggregation_weight_map);
         aggregation_weight_map.value() = 1.0;
       }
+      if (weightedrankmap.valid()) {
+        assign_pos_of(dwi, 0, 3).to(weightedrankmap);
+        weightedrankmap.value() = r - threshold.cutoff_p;
+      }
       break;
     default: {
       if (m <= n)
@@ -727,6 +740,10 @@ template <typename F> class DenoisingFunctor {
         }
         out.row(3) += weight * X.col(voxel_index);
         aggregation_weight_map.value() += weight;
+        if (weightedrankmap.valid()) {
+          assign_pos_of(neighbourhood.voxels[voxel_index].offset, 0, 3).to(weightedrankmap);
+          weightedrankmap.value() += weight * (r - threshold.cutoff_p);
+        }
       }
     } break;
     }
@@ -752,7 +769,7 @@ template <typename F> class DenoisingFunctor {
       assign_pos_of(dwi, 0, 3).to(voxelsmap);
       voxelsmap.value() = n;
     }
-  }
+  } // End functor
 
 private:
   // Denoising configuration
@@ -784,6 +801,7 @@ template <typename F> class DenoisingFunctor {
   // TODO Group these into a class?
   Image<float> noise;
   Image<uint16_t> rankmap;
+  Image<float> weightedrankmap;
   Image<float> sumweightsmap;
   Image<float> maxdistmap;
   Image<uint16_t> voxelsmap;
@@ -810,6 +828,7 @@ void run(Header &data,
          Image<bool> &mask,
          Image<float> &noise,
          Image<uint16_t> &rank,
+         Image<float> &weighted_rank,
          Image<float> &sum_weights,
          Image<float> &max_dist,
          Image<uint16_t> &voxels,
@@ -825,8 +844,19 @@ void run(Header &data,
   header.datatype() = DataType::from<T>();
   auto output = Image<T>::create(output_name, header);
   // run
-  DenoisingFunctor<T> func(
-      data, kernel, filter, aggregator, mask, noise, rank, sum_weights, max_dist, voxels, aggregation_sum, estimator);
+  DenoisingFunctor<T> func(data,
+                           kernel,
+                           filter,
+                           aggregator,
+                           mask,
+                           noise,
+                           rank,
+                           weighted_rank,
+                           sum_weights,
+                           max_dist,
+                           voxels,
+                           aggregation_sum,
+                           estimator);
   ThreadedLoop("running MP-PCA denoising", data, 0, 3).run(func, input, output);
   // Rescale output if performing aggregation
   if (aggregator == aggregator_type::EXCLUSIVE)
@@ -835,6 +865,10 @@ void run(Header &data,
     for (auto l_volume = Loop(3)(output); l_volume; ++l_volume)
       output.value() /= float(aggregation_sum.value());
   }
+  if (weighted_rank.valid()) {
+    for (auto l = Loop(aggregation_sum)(weighted_rank, aggregation_sum); l; ++l)
+      weighted_rank.value() /= aggregation_sum.value();
+  }
 }
 
 void run() {
@@ -877,33 +911,47 @@ void run() {
   if (!opt.empty())
     aggregator = aggregator_type(int(opt[0][0]));
 
+  Header H3D(dwi);
+  H3D.ndim() = 3;
+  H3D.reset_intensity_scaling();
+
   Image<float> noise;
   opt = get_options("noise");
   if (!opt.empty()) {
-    Header header(dwi);
-    header.ndim() = 3;
+    Header header(H3D);
     header.datatype() = DataType::Float32;
+    header.datatype().set_byte_order_native();
     noise = Image<float>::create(opt[0][0], header);
   }
 
   Image<uint16_t> rank;
   opt = get_options("rank");
   if (!opt.empty()) {
-    Header header(dwi);
-    header.ndim() = 3;
+    Header header(H3D);
     header.datatype() = DataType::UInt16;
-    header.reset_intensity_scaling();
     rank = Image<uint16_t>::create(opt[0][0], header);
   }
 
+  Image<float> weighted_rank;
+  opt = get_options("weightedrank");
+  if (!opt.empty()) {
+    if (aggregator == aggregator_type::EXCLUSIVE) {
+      WARN("When using -aggregator exclusive, "
+           "the output of -weightedrank will be identical to the output of -rank, "
+           "as there is no aggregation of multiple patches per output voxel");
+    }
+    Header header(H3D);
+    header.datatype() = DataType::Float32;
+    header.datatype().set_byte_order_native();
+    weighted_rank = Image<float>::create(opt[0][0], header);
+  }
+
   Image<float> sum_weights;
   opt = get_options("sumweights");
   if (!opt.empty()) {
-    Header header(dwi);
-    header.ndim() = 3;
+    Header header(H3D);
     header.datatype() = DataType::Float32;
     header.datatype().set_byte_order_native();
-    header.reset_intensity_scaling();
     sum_weights = Image<float>::create(opt[0][0], header);
     if (filter == filter_type::TRUNCATE) {
       WARN("Note that with a truncation filter, "
@@ -914,41 +962,31 @@ void run() {
   Image<float> max_dist;
   opt = get_options("max_dist");
   if (!opt.empty()) {
-    Header header(dwi);
-    header.ndim() = 3;
+    Header header(H3D);
     header.datatype() = DataType::Float32;
     header.datatype().set_byte_order_native();
-    header.reset_intensity_scaling();
     max_dist = Image<float>::create(opt[0][0], header);
   }
 
   Image<uint16_t> voxels;
   opt = get_options("voxels");
   if (!opt.empty()) {
-    Header header(dwi);
-    header.ndim() = 3;
+    Header header(H3D);
     header.datatype() = DataType::UInt16;
-    header.reset_intensity_scaling();
+    header.datatype().set_byte_order_native();
     voxels = Image<uint16_t>::create(opt[0][0], header);
   }
 
   Image<float> aggregation_sum;
-  Header header_aggregation(dwi);
-  header_aggregation.ndim() = 3;
-  header_aggregation.datatype() = DataType::Float64;
+  Header header_aggregation(H3D);
+  header_aggregation.datatype() = DataType::Float32;
   header_aggregation.datatype().set_byte_order_native();
-  header_aggregation.reset_intensity_scaling();
   opt = get_options("aggregation_sum");
   if (!opt.empty()) {
     if (aggregator == aggregator_type::EXCLUSIVE) {
       WARN("Output from -aggregation_sum will just contain 1 for every voxel processed: "
            "no patch aggregation takes place when output series comex exclusively from central patch");
     }
-    Header header(dwi);
-    header.ndim() = 3;
-    header.datatype() = DataType::Float32;
-    header.datatype().set_byte_order_native();
-    header.reset_intensity_scaling();
     aggregation_sum = Image<float>::create(opt[0][0], header_aggregation);
   } else if (aggregator != aggregator_type::EXCLUSIVE) {
     aggregation_sum = Image<float>::scratch(header_aggregation, "Scratch buffer for patch aggregation weights");
@@ -1019,6 +1057,7 @@ void run() {
                mask,
                noise,
                rank,
+               weighted_rank,
                sum_weights,
                max_dist,
                voxels,
@@ -1035,6 +1074,7 @@ void run() {
                 mask,
                 noise,
                 rank,
+                weighted_rank,
                 sum_weights,
                 max_dist,
                 voxels,
@@ -1051,6 +1091,7 @@ void run() {
                 mask,
                 noise,
                 rank,
+                weighted_rank,
                 sum_weights,
                 max_dist,
                 voxels,
@@ -1067,6 +1108,7 @@ void run() {
                  mask,
                  noise,
                  rank,
+                 weighted_rank,
                  sum_weights,
                  max_dist,
                  voxels,

From 2e6b024cd9e17a79133a3455667029541e481f5f Mon Sep 17 00:00:00 2001
From: Robert Smith <robert.smith@florey.edu.au>
Date: Sat, 16 Nov 2024 23:56:32 +1100
Subject: [PATCH 17/34] dwidenoise: Bulk move code to src/

---
 cmd/dwidenoise.cpp                   | 983 +++------------------------
 src/denoise/denoise.cpp              |  30 +
 src/denoise/denoise.h                |  37 +
 src/denoise/estimator/base.h         |  30 +
 src/denoise/estimator/estimator.cpp  |  54 ++
 src/denoise/estimator/estimator.h    |  34 +
 src/denoise/estimator/exp.h          |  62 ++
 src/denoise/estimator/mrm2022.h      |  58 ++
 src/denoise/estimator/result.h       |  28 +
 src/denoise/exports.h                |  57 ++
 src/denoise/functor.cpp              | 230 +++++++
 src/denoise/functor.h                |  91 +++
 src/denoise/kernel/base.h            |  39 ++
 src/denoise/kernel/cuboid.cpp        |  67 ++
 src/denoise/kernel/cuboid.h          |  42 ++
 src/denoise/kernel/data.h            |  35 +
 src/denoise/kernel/kernel.cpp        | 128 ++++
 src/denoise/kernel/kernel.h          |  39 ++
 src/denoise/kernel/sphere_base.cpp   |  45 ++
 src/denoise/kernel/sphere_base.h     |  54 ++
 src/denoise/kernel/sphere_radius.cpp |  37 +
 src/denoise/kernel/sphere_radius.h   |  42 ++
 src/denoise/kernel/sphere_ratio.cpp  |  62 ++
 src/denoise/kernel/sphere_ratio.h    |  51 ++
 src/denoise/kernel/voxel.h           |  54 ++
 25 files changed, 1516 insertions(+), 873 deletions(-)
 create mode 100644 src/denoise/denoise.cpp
 create mode 100644 src/denoise/denoise.h
 create mode 100644 src/denoise/estimator/base.h
 create mode 100644 src/denoise/estimator/estimator.cpp
 create mode 100644 src/denoise/estimator/estimator.h
 create mode 100644 src/denoise/estimator/exp.h
 create mode 100644 src/denoise/estimator/mrm2022.h
 create mode 100644 src/denoise/estimator/result.h
 create mode 100644 src/denoise/exports.h
 create mode 100644 src/denoise/functor.cpp
 create mode 100644 src/denoise/functor.h
 create mode 100644 src/denoise/kernel/base.h
 create mode 100644 src/denoise/kernel/cuboid.cpp
 create mode 100644 src/denoise/kernel/cuboid.h
 create mode 100644 src/denoise/kernel/data.h
 create mode 100644 src/denoise/kernel/kernel.cpp
 create mode 100644 src/denoise/kernel/kernel.h
 create mode 100644 src/denoise/kernel/sphere_base.cpp
 create mode 100644 src/denoise/kernel/sphere_base.h
 create mode 100644 src/denoise/kernel/sphere_radius.cpp
 create mode 100644 src/denoise/kernel/sphere_radius.h
 create mode 100644 src/denoise/kernel/sphere_ratio.cpp
 create mode 100644 src/denoise/kernel/sphere_ratio.h
 create mode 100644 src/denoise/kernel/voxel.h

diff --git a/cmd/dwidenoise.cpp b/cmd/dwidenoise.cpp
index a99bfd8806..d90b0e1c7c 100644
--- a/cmd/dwidenoise.cpp
+++ b/cmd/dwidenoise.cpp
@@ -15,29 +15,33 @@
  */
 
 #include <mutex>
+#include <string>
+#include <vector>
 
 #include "command.h"
+#include "header.h"
 #include "image.h"
 
 #include <Eigen/Dense>
 #include <Eigen/Eigenvalues>
 
+#include "denoise/denoise.h"
+#include "denoise/estimator/base.h"
+#include "denoise/estimator/estimator.h"
+#include "denoise/estimator/exp.h"
+#include "denoise/estimator/mrm2022.h"
+#include "denoise/estimator/result.h"
+#include "denoise/exports.h"
+#include "denoise/functor.h"
+#include "denoise/kernel/cuboid.h"
+#include "denoise/kernel/data.h"
+#include "denoise/kernel/kernel.h"
+#include "denoise/kernel/sphere_radius.h"
+#include "denoise/kernel/sphere_ratio.h"
+
 using namespace MR;
 using namespace App;
-
-const std::vector<std::string> dtypes = {"float32", "float64"};
-const std::vector<std::string> estimators = {"exp1", "exp2", "mrm2022"};
-enum class estimator_type { EXP1, EXP2, MRM2022 };
-
-const std::vector<std::string> shapes = {"cuboid", "sphere"};
-enum class shape_type { CUBOID, SPHERE };
-constexpr default_type sphere_multiplier_default = 1.0 / 0.85;
-
-const std::vector<std::string> filters = {"truncate", "frobenius"};
-enum class filter_type { TRUNCATE, FROBENIUS };
-
-const std::vector<std::string> aggregators = {"exclusive", "gaussian", "invl0", "rank", "uniform"};
-enum class aggregator_type { EXCLUSIVE, GAUSSIAN, INVL0, RANK, UNIFORM };
+using namespace MR::Denoise;
 
 // clang-format off
 void usage() {
@@ -62,27 +66,9 @@ void usage() {
     " If available, including the MRI phase data can reduce such non-Gaussian biases,"
     " and the command now supports complex input data."
 
-  + "The sliding spatial window behaves differently at the edges of the image FoV "
-    "depending on the shape / size selected for that window. "
-    "The default behaviour is to use a spherical kernel centred at the voxel of interest, "
-    "whose size is some multiple of the number of input volumes; "
-    "where some such voxels lie outside of the image FoV, "
-    "the radius of the kernel will be increased until the requisite number of voxels are used. "
-    "For a spherical kernel of a fixed radius, "
-    "no such expansion will occur, "
-    "and so for voxels near the image edge a reduced number of voxels will be present in the kernel. "
-    "For a cuboid kernel, "
-    "the centre of the kernel will be offset from the voxel being processed "
-    "such that the entire volume of the kernel resides within the image FoV."
+  + Kernel::shape_description
 
-  + "The size of the default spherical kernel is set to select a number of voxels that is "
-    "1.1 times the number of volumes in the input series. "
-    "If a cuboid kernel is requested, "
-    "but the -extent option is not specified, "
-    "the command will select the smallest isotropic patch size "
-    "that exceeds the number of DW images in the input data; "
-    "e.g., 5x5x5 for data with <= 125 DWI volumes, "
-    "7x7x7 for data with <= 343 DWI volumes, etc."
+  + Kernel::size_description
 
   + "By default, optimal value shrinkage based on minimisation of the Frobenius norm "
     "will be used to attenuate eigenvectors based on the estimated noise level. "
@@ -141,27 +127,21 @@ void usage() {
   + Argument("out", "the output denoised DWI image.").type_image_out();
 
   OPTIONS
-  + OptionGroup("Options for modifying the application of PCA denoising")
+  + OptionGroup("Options for modifying PCA computations")
+  + datatype_option
+  + Estimator::option
+
+  + Kernel::options
+
+  + OptionGroup("Options that affect reconstruction of the output image series")
   + Option("mask",
-           "Only process voxels within the specified binary brain mask image.")
+           "Only denoise voxels within the specified binary brain mask image.")
     + Argument("image").type_image_in()
-  + Option("datatype",
-           "Datatype for the eigenvalue decomposition"
-           " (single or double precision). "
-           "For complex input data,"
-           " this will select complex float32 or complex float64 datatypes.")
-    + Argument("float32/float64").type_choice(dtypes)
-  + Option("estimator",
-           "Select the noise level estimator"
-           " (default = Exp2),"
-           " either: \n"
-           "* Exp1: the original estimator used in Veraart et al. (2016); \n"
-           "* Exp2: the improved estimator introduced in Cordero-Grande et al. (2019); \n"
-           "* MRM2022: the alternative estimator introduced in Olesen et al. (2022).")
-    + Argument("algorithm").type_choice(estimators)
   + Option("filter",
-           "Modulate how components are filtered based on their eigenvalues; "
-           "options are: " + join(filters, ",") + "; default: frobenius")
+           "Modulate how component contributions are filtered "
+           "based on the cumulative eigenvalues relative to the noise level; "
+           "options are: " + join(filters, ",") + "; "
+           "default: frobenius (Optimal Shrinkage based on minimisation of the Frobenius norm)")
     + Argument("choice").type_choice(filters)
   + Option("aggregator",
            "Select how the outcomes of multiple PCA outcomes centred at different voxels "
@@ -171,54 +151,54 @@ void usage() {
   // TODO For specifically the Gaussian aggregator,
   //   should ideally be possible to select the FWHM of the aggregator
 
-  // TODO Consider renaming some options to better distinguish between:
-  // - Parameters arising from PCA-based noise level estimation
-  // - Parameters encoding properties of the output data
-
+  // TODO Consider restructuring & encapsulating in classes
+  // Also bear in mind how use of subsampling could affect this
+  // TODO If using downsampling,
+  //   may be preferable to explicitly downsample the voxel grid on which to yield these data
+  // - -noise: This is a patch-wise estimate;
+  //           could however elect to make it an aggreate mean if aggregation is being performed
+  //           (or just forbid it if downsampling is performed?)
+  // - -rank: This is a patch-wise estimate (see -noise above)
+  // - -weightedrank: This is explicitly a multi-patch aggregation,
+  //   so doesn't apply to noise level estimation only;
+  //   it may be better to distinguish between "input rank" and "output rank"?
+  // - -sumweights: Explicitly a multi-patch aggregation
+  // - -max_dist: This is a property of the local patch;
+  //              this could be hidden behind a #define TBH
+  // - -voxels: This is a property of the local patch (see -max_dist above)
+  // - -aggregation_sum: This is explicitly a multi-patch aggregation
+  // TODO Consider an option group for "debugging of sliding window kernel behaviour"
+  // TODO Potential issue wherein optimal shrinkage may set to 0
+  //   some components above the determined noise level...
   + OptionGroup("Options for exporting additional data regarding PCA behaviour")
-  + Option("noise",
+  + Option("noise_out",
            "The output noise map,"
            " i.e., the estimated noise level 'sigma' in the data. "
            "Note that on complex input data,"
            " this will be the total noise level across real and imaginary channels,"
            " so a scale factor sqrt(2) applies.")
     + Argument("image").type_image_out()
-  + Option("rank",
-           "The estimated signal rank for the denoising patch centred at each voxel")
+  + Option("rank_input",
+           "The signal rank estimated for the denoising patch centred at each input image voxel")
     + Argument("image").type_image_out()
-  + Option("weightedrank",
-           "The weighted mean rank for the output image data, accounting for multi-patch aggregation")
-    + Argument("image").type_image_out()
-  + Option("sumweights",
-           "the sum of eigenvector weights computed for the denoising patch centred at each voxel")
+  + Option("rank_output",
+           "An estimated rank for the output image data, accounting for multi-patch aggregation")
     + Argument("image").type_image_out()
+
+  + OptionGroup("Options for debugging the operation of sliding window kernels")
   + Option("max_dist",
            "The maximum distance between a voxel and another voxel that was included in the local denoising patch")
     + Argument("image").type_image_out()
-  + Option("voxels",
+  + Option("voxelcount",
            "The number of voxels that contributed to the PCA for processing of each voxel")
     + Argument("image").type_image_out()
-  + Option("aggregation_sum",
+  + Option("sum_aggregation",
            "The sum of aggregation weights of those patches contributing to each output voxel")
     + Argument("image").type_image_out()
-
-  + OptionGroup("Options for controlling the sliding spatial window")
-  + Option("shape",
-           "Set the shape of the sliding spatial window. "
-           "Options are: " + join(shapes, ",") + "; default: sphere")
-    + Argument("choice").type_choice(shapes)
-  + Option("radius_mm",
-           "Set an absolute spherical kernel radius in mm")
-    + Argument("value").type_float(0.0)
-  + Option("radius_ratio",
-           "Set the spherical kernel size as a ratio of number of voxels to number of input volumes "
-           "(default: 1.0/0.85 ~= 1.18)")
-    + Argument("value").type_float(0.0)
-  // TODO Command-line option that allows user to specify minimum absolute number of voxels in kernel
-  + Option("extent",
-           "Set the patch size of the cuboid kernel; "
-           "can be either a single odd integer or a comma-separated triplet of odd integers")
-    + Argument("window").type_sequence_int();
+  + Option("sum_optshrink",
+           "the sum of eigenvector weights computed for the denoising patch centred at each voxel "
+           "as a result of performing optimal shrinkage")
+    + Argument("image").type_image_out();
 
   COPYRIGHT =
       "Copyright (c) 2016 New York University, University of Antwerp, and the MRtrix3 contributors \n \n"
@@ -249,575 +229,6 @@ void usage() {
 }
 // clang-format on
 
-using voxel_type = Eigen::Array<int, 3, 1>;
-using vector_type = Eigen::VectorXd;
-
-class KernelVoxel {
-public:
-  KernelVoxel(const voxel_type &offset, const default_type sq_distance) : offset(offset), sq_distance(sq_distance) {}
-  KernelVoxel(const KernelVoxel &) = default;
-  KernelVoxel(KernelVoxel &&) = default;
-  ~KernelVoxel() {}
-  KernelVoxel &operator=(const KernelVoxel &that) {
-    offset = that.offset;
-    sq_distance = that.sq_distance;
-    return *this;
-  }
-  KernelVoxel &operator=(KernelVoxel &&that) noexcept {
-    offset = that.offset;
-    sq_distance = that.sq_distance;
-    return *this;
-  }
-  bool operator<(const KernelVoxel &that) const { return sq_distance < that.sq_distance; }
-  default_type distance() const { return std::sqrt(sq_distance); }
-  // TODO Sometimes this acts as an offset, other times it acts as an absolute voxel index
-  // Consider either renaming, or actually using two different classes
-  // The latter could use ssize_t instead of int to better indicate this
-  voxel_type offset;
-  default_type sq_distance;
-};
-
-// Class to encode return information from kernel
-class KernelData {
-public:
-  KernelData() : centre_index(-1), max_distance(-std::numeric_limits<default_type>::infinity()) {}
-  KernelData(const ssize_t i) : centre_index(i), max_distance(-std::numeric_limits<default_type>::infinity()) {}
-  std::vector<KernelVoxel> voxels;
-  ssize_t centre_index;
-  default_type max_distance;
-};
-
-class KernelBase {
-public:
-  KernelBase(const Header &H) : H(H) {}
-  KernelBase(const KernelBase &) = default;
-  virtual ~KernelBase() = default;
-  // This is just for pre-allocating matrices
-  virtual ssize_t estimated_size() const = 0;
-  // This is the interface that kernels must provide
-  virtual KernelData operator()(const voxel_type &) const = 0;
-
-protected:
-  const Header H;
-};
-
-class KernelCube : public KernelBase {
-public:
-  KernelCube(const Header &header, const std::vector<uint32_t> &extent)
-      : KernelBase(header),
-        half_extent({int(extent[0] / 2), int(extent[1] / 2), int(extent[2] / 2)}),
-        size(ssize_t(extent[0]) * ssize_t(extent[1]) * ssize_t(extent[2])),
-        centre_index(size / 2) {
-    for (auto e : extent) {
-      if (!(e % 2))
-        throw Exception("Size of cubic kernel must be an odd integer");
-    }
-  }
-  KernelCube(const KernelCube &) = default;
-  ~KernelCube() final = default;
-  KernelData operator()(const voxel_type &pos) const override {
-    KernelData result(centre_index);
-    voxel_type voxel;
-    voxel_type offset;
-    for (offset[2] = -half_extent[2]; offset[2] <= half_extent[2]; ++offset[2]) {
-      voxel[2] = wrapindex(pos[2], offset[2], half_extent[2], H.size(2));
-      for (offset[1] = -half_extent[1]; offset[1] <= half_extent[1]; ++offset[1]) {
-        voxel[1] = wrapindex(pos[1], offset[1], half_extent[1], H.size(1));
-        for (offset[0] = -half_extent[0]; offset[0] <= half_extent[0]; ++offset[0]) {
-          voxel[0] = wrapindex(pos[0], offset[0], half_extent[0], H.size(0));
-          const default_type sq_distance = Math::pow2((pos[0] - voxel[0]) * H.spacing(0)) +
-                                           Math::pow2((pos[1] - voxel[1]) * H.spacing(1)) +
-                                           Math::pow2((pos[2] - voxel[2]) * H.spacing(2));
-          result.voxels.push_back(KernelVoxel(voxel, sq_distance));
-          result.max_distance = std::max(result.max_distance, sq_distance);
-        }
-      }
-    }
-    result.max_distance = std::sqrt(result.max_distance);
-    return result;
-  }
-  ssize_t estimated_size() const override { return size; }
-
-private:
-  const std::vector<int> dimensions;
-  const std::vector<int> half_extent;
-  const ssize_t size;
-  const ssize_t centre_index;
-
-  // patch handling at image edges
-  inline size_t wrapindex(int p, int r, int e, int max) const {
-    int rr = p + r;
-    if (rr < 0)
-      rr = e - r;
-    if (rr >= max)
-      rr = (max - 1) - e - r;
-    return rr;
-  }
-};
-
-class KernelSphereBase : public KernelBase {
-public:
-  KernelSphereBase(const Header &voxel_grid, const default_type max_radius)
-      : KernelBase(voxel_grid), shared(new Shared(voxel_grid, max_radius)) {}
-  KernelSphereBase(const KernelSphereBase &) = default;
-  virtual ~KernelSphereBase() override {}
-
-protected:
-  class Shared {
-  public:
-    using TableType = std::vector<KernelVoxel>;
-    Shared(const Header &voxel_grid, const default_type max_radius) {
-      const default_type max_radius_sq = Math::pow2(max_radius);
-      const voxel_type half_extents({int(std::ceil(max_radius / voxel_grid.spacing(0))),   //
-                                     int(std::ceil(max_radius / voxel_grid.spacing(1))),   //
-                                     int(std::ceil(max_radius / voxel_grid.spacing(2)))}); //
-      // Build the searchlight
-      data.reserve(size_t(2 * half_extents[0] + 1) * size_t(2 * half_extents[1] + 1) * size_t(2 * half_extents[2] + 1));
-      voxel_type offset({-1, -1, -1});
-      for (offset[2] = -half_extents[2]; offset[2] <= half_extents[2]; ++offset[2]) {
-        for (offset[1] = -half_extents[1]; offset[1] <= half_extents[1]; ++offset[1]) {
-          for (offset[0] = -half_extents[0]; offset[0] <= half_extents[0]; ++offset[0]) {
-            const default_type squared_distance = Math::pow2(offset[0] * voxel_grid.spacing(0))    //
-                                                  + Math::pow2(offset[1] * voxel_grid.spacing(1))  //
-                                                  + Math::pow2(offset[2] * voxel_grid.spacing(2)); //
-            if (squared_distance <= max_radius_sq)
-              data.emplace_back(KernelVoxel(offset, squared_distance));
-          }
-        }
-      }
-      std::sort(data.begin(), data.end());
-    }
-    TableType::const_iterator begin() const { return data.begin(); }
-    TableType::const_iterator end() const { return data.end(); }
-
-  private:
-    TableType data;
-  };
-  std::shared_ptr<Shared> shared;
-};
-
-class KernelSphereRatio : public KernelSphereBase {
-public:
-  KernelSphereRatio(const Header &voxel_grid, const default_type min_ratio)
-      : KernelSphereBase(voxel_grid, compute_max_radius(voxel_grid, min_ratio)),
-        min_size(std::ceil(voxel_grid.size(3) * min_ratio)) {}
-  KernelSphereRatio(const KernelSphereRatio &) = default;
-  ~KernelSphereRatio() final = default;
-  KernelData operator()(const voxel_type &pos) const override {
-    KernelData result(0);
-    auto table_it = shared->begin();
-    while (table_it != shared->end()) {
-      // If there's a tie in distances, want to include all such offsets in the kernel,
-      //   even if the size of the utilised kernel extends beyond the minimum size
-      if (result.voxels.size() >= min_size && table_it->sq_distance != result.max_distance)
-        break;
-      const voxel_type voxel({pos[0] + table_it->offset[0],   //
-                              pos[1] + table_it->offset[1],   //
-                              pos[2] + table_it->offset[2]}); //
-      if (!is_out_of_bounds(H, voxel, 0, 3)) {
-        result.voxels.push_back(KernelVoxel(voxel, table_it->sq_distance));
-        result.max_distance = table_it->sq_distance;
-      }
-      ++table_it;
-    }
-    if (table_it == shared->end()) {
-      throw Exception(                                                                   //
-          std::string("Inadequate spherical kernel initialisation ")                     //
-          + "(lookup table " + str(std::distance(shared->begin(), shared->end())) + "; " //
-          + "min size " + str(min_size) + "; "                                           //
-          + "read size " + str(result.voxels.size()) + ")");                             //
-    }
-    result.max_distance = std::sqrt(result.max_distance);
-    return result;
-  }
-  ssize_t estimated_size() const override { return min_size; }
-
-private:
-  ssize_t min_size;
-  // Determine an appropriate bounding box from which to generate the search table
-  // Find the radius for which 7/8 of the sphere will contain the minimum number of voxels, then round up
-  // This is only for setting the maximal radius for generation of the lookup table
-  default_type compute_max_radius(const Header &voxel_grid, const default_type min_ratio) const {
-    const size_t num_volumes = voxel_grid.size(3);
-    const default_type voxel_volume = voxel_grid.spacing(0) * voxel_grid.spacing(1) * voxel_grid.spacing(2);
-    const default_type sphere_volume = 8.0 * num_volumes * min_ratio * voxel_volume;
-    const default_type approx_radius = std::sqrt(sphere_volume * 0.75 / Math::pi);
-    const voxel_type half_extents({int(std::ceil(approx_radius / voxel_grid.spacing(0))),   //
-                                   int(std::ceil(approx_radius / voxel_grid.spacing(1))),   //
-                                   int(std::ceil(approx_radius / voxel_grid.spacing(2)))}); //
-    return std::max({half_extents[0] * voxel_grid.spacing(0),
-                     half_extents[1] * voxel_grid.spacing(1),
-                     half_extents[2] * voxel_grid.spacing(2)});
-  }
-};
-
-class KernelSphereFixedRadius : public KernelSphereBase {
-public:
-  KernelSphereFixedRadius(const Header &voxel_grid, const default_type radius)
-      : KernelSphereBase(voxel_grid, radius),                         //
-        maximum_size(std::distance(shared->begin(), shared->end())) { //
-    INFO("Maximum number of voxels in " + str(radius) + "mm fixed-radius kernel is " + str(maximum_size));
-  }
-  KernelSphereFixedRadius(const KernelSphereFixedRadius &) = default;
-  ~KernelSphereFixedRadius() final = default;
-  KernelData operator()(const voxel_type &pos) const {
-    KernelData result(0);
-    result.voxels.reserve(maximum_size);
-    for (auto map_it = shared->begin(); map_it != shared->end(); ++map_it) {
-      const voxel_type voxel({pos[0] + map_it->offset[0],   //
-                              pos[1] + map_it->offset[1],   //
-                              pos[2] + map_it->offset[2]}); //
-      if (!is_out_of_bounds(H, voxel, 0, 3)) {
-        result.voxels.push_back(KernelVoxel(voxel, map_it->sq_distance));
-        result.max_distance = map_it->sq_distance;
-      }
-    }
-    result.max_distance = std::sqrt(result.max_distance);
-    return result;
-  }
-  ssize_t estimated_size() const override { return maximum_size; }
-
-private:
-  const ssize_t maximum_size;
-};
-
-class EstimatorResult {
-public:
-  EstimatorResult() : cutoff_p(0), sigma2(0.0) {}
-  ssize_t cutoff_p;
-  double sigma2;
-};
-
-class EstimatorBase {
-public:
-  EstimatorBase() = default;
-  virtual EstimatorResult operator()(const vector_type &eigenvalues, const ssize_t m, const ssize_t n) const = 0;
-};
-
-template <ssize_t version> class EstimatorExp : public EstimatorBase {
-public:
-  EstimatorExp() = default;
-  EstimatorResult operator()(const vector_type &s, const ssize_t m, const ssize_t n) const final {
-    EstimatorResult result;
-    const ssize_t r = std::min(m, n);
-    const ssize_t q = std::max(m, n);
-    const double lam_r = std::max(s[0], 0.0) / q;
-    double clam = 0.0;
-    for (ssize_t p = 0; p < r; ++p) // p+1 is the number of noise components
-    {                               // (as opposed to the paper where p is defined as the number of signal components)
-      const double lam = std::max(s[p], 0.0) / q;
-      clam += lam;
-      double denominator = std::numeric_limits<double>::signaling_NaN();
-      switch (version) {
-      case 1:
-        denominator = q;
-        break;
-      case 2:
-        denominator = q - (r - p - 1);
-        break;
-      default:
-        assert(false);
-      }
-      const double gam = double(p + 1) / denominator;
-      const double sigsq1 = clam / double(p + 1);
-      const double sigsq2 = (lam - lam_r) / (4.0 * std::sqrt(gam));
-      // sigsq2 > sigsq1 if signal else noise
-      if (sigsq2 < sigsq1) {
-        result.sigma2 = sigsq1;
-        result.cutoff_p = p + 1;
-      }
-    }
-    return result;
-  }
-};
-
-class EstimatorMRM2022 : public EstimatorBase {
-public:
-  EstimatorMRM2022() = default;
-  EstimatorResult operator()(const vector_type &s, const ssize_t m, const ssize_t n) const final {
-    EstimatorResult result;
-    const ssize_t mprime = std::min(m, n);
-    const ssize_t nprime = std::max(m, n);
-    const double sigmasq_to_lamplus = Math::pow2(std::sqrt(nprime) + std::sqrt(mprime));
-    double clam = 0.0;
-    for (ssize_t i = 0; i != mprime; ++i)
-      clam += std::max(s[i], 0.0);
-    clam /= nprime;
-    // Unlike Exp# code,
-    //   MRM2022 article uses p to index number of signal components,
-    //   and here doing a direct translation of the manuscript content to code
-    double lamplusprev = -std::numeric_limits<double>::infinity();
-    for (ssize_t p = 0; p < mprime; ++p) {
-      const ssize_t i = mprime - 1 - p;
-      const double lam = std::max(s[i], 0.0) / nprime;
-      if (lam < lamplusprev)
-        return result;
-      clam -= lam;
-      const double sigmasq = clam / ((mprime - p) * (nprime - p));
-      lamplusprev = sigmasq * sigmasq_to_lamplus;
-      result.cutoff_p = i;
-      result.sigma2 = sigmasq;
-    }
-    return result;
-  }
-};
-
-template <typename F> class DenoisingFunctor {
-
-public:
-  using MatrixType = Eigen::Matrix<F, Eigen::Dynamic, Eigen::Dynamic>;
-
-  DenoisingFunctor(const Header &header,
-                   std::shared_ptr<KernelBase> kernel,
-                   filter_type filter,
-                   aggregator_type aggregator,
-                   Image<bool> &mask,
-                   Image<float> &noise,
-                   Image<uint16_t> &rank,
-                   Image<float> &weighted_rank,
-                   Image<float> &sum_weights,
-                   Image<float> &max_dist,
-                   Image<uint16_t> &voxels,
-                   // TODO Would be preferable for this to be double if computations are happening using double
-                   Image<float> &aggregation_weight_map,
-                   std::shared_ptr<EstimatorBase> estimator)
-      : kernel(kernel),
-        filter(filter),
-        aggregator(aggregator),
-        // FWHM = 2 x cube root of voxel spacings
-        gaussian_multiplier(-std::log(2.0) /
-                            Math::pow2(std::cbrt(header.spacing(0) * header.spacing(1) * header.spacing(2)))),
-        m(header.size(3)),
-        estimator(estimator),
-        mask(mask),
-        X(m, kernel->estimated_size()),
-        XtX(std::min(m, kernel->estimated_size()), std::min(m, kernel->estimated_size())),
-        eig(std::min(m, kernel->estimated_size())),
-        s(std::min(m, kernel->estimated_size())),
-        clam(std::min(m, kernel->estimated_size())),
-        w(std::min(m, kernel->estimated_size())),
-        noise(noise),
-        rankmap(rank),
-        weightedrankmap(weighted_rank),
-        sumweightsmap(sum_weights),
-        maxdistmap(max_dist),
-        voxelsmap(voxels),
-        aggregation_weight_map(aggregation_weight_map) {}
-
-  template <typename ImageType> void operator()(ImageType &dwi, ImageType &out) {
-    // Process voxels in mask only
-    if (mask.valid()) {
-      assign_pos_of(dwi, 0, 3).to(mask);
-      if (!mask.value())
-        return;
-    }
-
-    // Load list of voxels from which to load data
-    const KernelData neighbourhood = (*kernel)({int(dwi.index(0)), int(dwi.index(1)), int(dwi.index(2))});
-    const ssize_t n = neighbourhood.voxels.size();
-    const ssize_t r = std::min(m, n);
-    const ssize_t q = std::max(m, n);
-
-    // Expand local storage if necessary
-    if (n > X.cols()) {
-      DEBUG("Expanding data matrix storage from " + str(m) + "x" + str(X.cols()) + " to " + str(m) + "x" + str(n));
-      X.resize(m, n);
-    }
-    if (r > XtX.cols()) {
-      DEBUG("Expanding decomposition matrix storage from " + str(X.rows()) + " to " + str(r));
-      XtX.resize(r, r);
-      s.resize(r);
-      clam.resize(r);
-      w.resize(r);
-    }
-
-    // Fill matrices with NaN when in debug mode;
-    //   make sure results from one voxel are not creeping into another
-    //   due to use of block oberations to prevent memory re-allocation
-    //   in the presence of variation in kernel sizes
-#ifndef NDEBUG
-    X.fill(std::numeric_limits<F>::signaling_NaN());
-    XtX.fill(std::numeric_limits<F>::signaling_NaN());
-    s.fill(std::numeric_limits<default_type>::signaling_NaN());
-    clam.fill(std::numeric_limits<default_type>::signaling_NaN());
-    w.fill(std::numeric_limits<default_type>::signaling_NaN());
-#endif
-
-    load_data(dwi, neighbourhood.voxels);
-
-    // Compute Eigendecomposition:
-    if (m <= n)
-      XtX.topLeftCorner(r, r).template triangularView<Eigen::Lower>() = X.leftCols(n) * X.leftCols(n).adjoint();
-    else
-      XtX.topLeftCorner(r, r).template triangularView<Eigen::Lower>() = X.leftCols(n).adjoint() * X.leftCols(n);
-    eig.compute(XtX.topLeftCorner(r, r));
-    // eigenvalues sorted in increasing order:
-    s.head(r) = eig.eigenvalues().template cast<double>();
-
-    // Marchenko-Pastur optimal threshold determination
-    const EstimatorResult threshold = (*estimator)(s, m, n);
-
-    // Generate weights vector
-    double sum_weights = 0.0;
-    switch (filter) {
-    case filter_type::TRUNCATE:
-      w.head(threshold.cutoff_p).setZero();
-      w.segment(threshold.cutoff_p, r - threshold.cutoff_p).setOnes();
-      sum_weights = r - threshold.cutoff_p;
-      break;
-    case filter_type::FROBENIUS: {
-      const double beta = r / q;
-      const double transition = 1.0 + std::sqrt(beta);
-      double clam = 0.0;
-      for (ssize_t i = 0; i != r; ++i) {
-        const double lam = std::max(s[i], 0.0) / q;
-        clam += lam;
-        const double y = clam / (threshold.sigma2 * (i + 1));
-        const double nu = y > transition ? std::sqrt(Math::pow2(Math::pow2(y) - beta - 1.0) - (4.0 * beta)) / y : 0.0;
-        w[i] = nu / y;
-        sum_weights += w[i];
-      }
-    } break;
-    default:
-      assert(false);
-    }
-
-    // recombine data using only eigenvectors above threshold
-    // If only the data computed when this voxel was the centre of the patch
-    //   is to be used for synthesis of the output image,
-    //   then only that individual column needs to be reconstructed;
-    //   if however the result from this patch is to contribute to the synthesized image
-    //   for all voxels that were utilised within this patch,
-    //   then we need to instead compute the full projection
-    switch (aggregator) {
-    case aggregator_type::EXCLUSIVE:
-      if (m <= n)
-        X.col(neighbourhood.centre_index) =
-            eig.eigenvectors() *
-            (w.head(r).cast<F>().asDiagonal() * (eig.eigenvectors().adjoint() * X.col(neighbourhood.centre_index)));
-      else
-        X.col(neighbourhood.centre_index) =
-            X.leftCols(n) * (eig.eigenvectors() * (w.head(r).cast<F>().asDiagonal() *
-                                                   eig.eigenvectors().adjoint().col(neighbourhood.centre_index)));
-      assign_pos_of(dwi).to(out);
-      out.row(3) = X.col(neighbourhood.centre_index);
-      if (aggregation_weight_map.valid()) {
-        assign_pos_of(dwi, 0, 3).to(aggregation_weight_map);
-        aggregation_weight_map.value() = 1.0;
-      }
-      if (weightedrankmap.valid()) {
-        assign_pos_of(dwi, 0, 3).to(weightedrankmap);
-        weightedrankmap.value() = r - threshold.cutoff_p;
-      }
-      break;
-    default: {
-      if (m <= n)
-        X = eig.eigenvectors() * (w.head(r).cast<F>().asDiagonal() * (eig.eigenvectors().adjoint() * X));
-      else
-        X.leftCols(n) =
-            X.leftCols(n) * (eig.eigenvectors() * (w.head(r).cast<F>().asDiagonal() * eig.eigenvectors().adjoint()));
-      std::lock_guard<std::mutex> lock(mutex_aggregator);
-      for (size_t voxel_index = 0; voxel_index != neighbourhood.voxels.size(); ++voxel_index) {
-        assign_pos_of(neighbourhood.voxels[voxel_index].offset, 0, 3).to(out);
-        assign_pos_of(neighbourhood.voxels[voxel_index].offset).to(aggregation_weight_map);
-        double weight = std::numeric_limits<double>::signaling_NaN();
-        switch (aggregator) {
-        case aggregator_type::EXCLUSIVE:
-          assert(false);
-          break;
-        case aggregator_type::GAUSSIAN:
-          weight = std::exp(gaussian_multiplier * neighbourhood.voxels[voxel_index].sq_distance);
-          break;
-        case aggregator_type::INVL0:
-          weight = 1.0 / (1 + r - threshold.cutoff_p);
-          break;
-        case aggregator_type::RANK:
-          weight = r - threshold.cutoff_p;
-          break;
-        case aggregator_type::UNIFORM:
-          weight = 1.0;
-          break;
-        }
-        out.row(3) += weight * X.col(voxel_index);
-        aggregation_weight_map.value() += weight;
-        if (weightedrankmap.valid()) {
-          assign_pos_of(neighbourhood.voxels[voxel_index].offset, 0, 3).to(weightedrankmap);
-          weightedrankmap.value() += weight * (r - threshold.cutoff_p);
-        }
-      }
-    } break;
-    }
-
-    // Store additional output maps if requested
-    if (noise.valid()) {
-      assign_pos_of(dwi, 0, 3).to(noise);
-      noise.value() = float(std::sqrt(threshold.sigma2));
-    }
-    if (rankmap.valid()) {
-      assign_pos_of(dwi, 0, 3).to(rankmap);
-      rankmap.value() = uint16_t(r - threshold.cutoff_p);
-    }
-    if (sumweightsmap.valid()) {
-      assign_pos_of(dwi, 0, 3).to(sumweightsmap);
-      sumweightsmap.value() = sum_weights;
-    }
-    if (maxdistmap.valid()) {
-      assign_pos_of(dwi, 0, 3).to(maxdistmap);
-      maxdistmap.value() = neighbourhood.max_distance;
-    }
-    if (voxelsmap.valid()) {
-      assign_pos_of(dwi, 0, 3).to(voxelsmap);
-      voxelsmap.value() = n;
-    }
-  } // End functor
-
-private:
-  // Denoising configuration
-  std::shared_ptr<KernelBase> kernel;
-  filter_type filter;
-  aggregator_type aggregator;
-  double gaussian_multiplier;
-  const ssize_t m;
-  std::shared_ptr<EstimatorBase> estimator;
-  Image<bool> mask;
-
-  // Reusable memory
-  MatrixType X;
-  MatrixType XtX;
-  Eigen::SelfAdjointEigenSolver<MatrixType> eig;
-  vector_type s;
-  vector_type clam;
-  vector_type w;
-
-  // Data that can only be written in a thread-safe manner
-  // Note that this applies not just to this scratch buffer, but also the output image
-  //   (while it would be thread-safe to create a full copy of the output image for each thread
-  //   and combine them only at destruction time,
-  //   this runs the risk of becoming prohibitively large)
-  // Not placing this within a MutexProtexted<> as the image type is still templated
-  static std::mutex mutex_aggregator;
-
-  // Export images
-  // TODO Group these into a class?
-  Image<float> noise;
-  Image<uint16_t> rankmap;
-  Image<float> weightedrankmap;
-  Image<float> sumweightsmap;
-  Image<float> maxdistmap;
-  Image<uint16_t> voxelsmap;
-  Image<float> aggregation_weight_map;
-
-  template <typename ImageType> void load_data(ImageType &image, const std::vector<KernelVoxel> &voxels) {
-    const voxel_type pos({int(image.index(0)), int(image.index(1)), int(image.index(2))});
-    for (ssize_t i = 0; i != voxels.size(); ++i) {
-      assign_pos_of(voxels[i].offset, 0, 3).to(image);
-      X.col(i) = image.row(3);
-    }
-    assign_pos_of(pos, 0, 3).to(image);
-  }
-};
-template <typename F> std::mutex DenoisingFunctor<F>::mutex_aggregator;
-
 // Necessary to allow normalisation by sum of aggregation weights
 //   where the image type is cdouble, but aggregation weights are float
 // (operations combining complex & real types not allowed to be of different precision)
@@ -826,48 +237,30 @@ std::complex<double> operator/(const std::complex<double> &c, const float n) { r
 template <typename T>
 void run(Header &data,
          Image<bool> &mask,
-         Image<float> &noise,
-         Image<uint16_t> &rank,
-         Image<float> &weighted_rank,
-         Image<float> &sum_weights,
-         Image<float> &max_dist,
-         Image<uint16_t> &voxels,
-         Image<float> &aggregation_sum,
-         const std::string &output_name,
-         std::shared_ptr<KernelBase> kernel,
+         std::shared_ptr<Kernel::Base> kernel,
+         std::shared_ptr<Estimator::Base> estimator,
          filter_type filter,
          aggregator_type aggregator,
-         std::shared_ptr<EstimatorBase> estimator) {
+         const std::string &output_name,
+         Exports &exports) {
   auto input = data.get_image<T>().with_direct_io(3);
   // create output
   Header header(data);
   header.datatype() = DataType::from<T>();
   auto output = Image<T>::create(output_name, header);
   // run
-  DenoisingFunctor<T> func(data,
-                           kernel,
-                           filter,
-                           aggregator,
-                           mask,
-                           noise,
-                           rank,
-                           weighted_rank,
-                           sum_weights,
-                           max_dist,
-                           voxels,
-                           aggregation_sum,
-                           estimator);
+  Functor<T> func(data, mask, kernel, estimator, filter, aggregator, exports);
   ThreadedLoop("running MP-PCA denoising", data, 0, 3).run(func, input, output);
   // Rescale output if performing aggregation
   if (aggregator == aggregator_type::EXCLUSIVE)
     return;
-  for (auto l_voxel = Loop(aggregation_sum)(output, aggregation_sum); l_voxel; ++l_voxel) {
+  for (auto l_voxel = Loop(exports.sum_aggregation)(output, exports.sum_aggregation); l_voxel; ++l_voxel) {
     for (auto l_volume = Loop(3)(output); l_volume; ++l_volume)
-      output.value() /= float(aggregation_sum.value());
+      output.value() /= float(exports.sum_aggregation.value());
   }
-  if (weighted_rank.valid()) {
-    for (auto l = Loop(aggregation_sum)(weighted_rank, aggregation_sum); l; ++l)
-      weighted_rank.value() /= aggregation_sum.value();
+  if (exports.rank_output.valid()) {
+    for (auto l = Loop(exports.sum_aggregation)(exports.rank_output, exports.sum_aggregation); l; ++l)
+      exports.rank_output.value() /= exports.sum_aggregation.value();
   }
 }
 
@@ -884,22 +277,11 @@ void run() {
     check_dimensions(mask, dwi, 0, 3);
   }
 
-  std::shared_ptr<EstimatorBase> estimator;
-  opt = get_options("estimator");
-  const estimator_type est = opt.empty() ? estimator_type::EXP2 : estimator_type((int)(opt[0][0]));
-  switch (est) {
-  case estimator_type::EXP1:
-    estimator = std::make_shared<EstimatorExp<1>>();
-    break;
-  case estimator_type::EXP2:
-    estimator = std::make_shared<EstimatorExp<2>>();
-    break;
-  case estimator_type::MRM2022:
-    estimator = std::make_shared<EstimatorMRM2022>();
-    break;
-  default:
-    assert(false);
-  }
+  auto kernel = Kernel::make_kernel(dwi);
+  assert(kernel);
+
+  auto estimator = Estimator::make_estimator();
+  assert(estimator);
 
   filter_type filter = filter_type::FROBENIUS;
   opt = get_options("filter");
@@ -911,141 +293,48 @@ void run() {
   if (!opt.empty())
     aggregator = aggregator_type(int(opt[0][0]));
 
-  Header H3D(dwi);
-  H3D.ndim() = 3;
-  H3D.reset_intensity_scaling();
-
-  Image<float> noise;
-  opt = get_options("noise");
-  if (!opt.empty()) {
-    Header header(H3D);
-    header.datatype() = DataType::Float32;
-    header.datatype().set_byte_order_native();
-    noise = Image<float>::create(opt[0][0], header);
-  }
-
-  Image<uint16_t> rank;
-  opt = get_options("rank");
-  if (!opt.empty()) {
-    Header header(H3D);
-    header.datatype() = DataType::UInt16;
-    rank = Image<uint16_t>::create(opt[0][0], header);
-  }
-
-  Image<float> weighted_rank;
-  opt = get_options("weightedrank");
+  Exports exports(dwi);
+  opt = get_options("noise_out");
+  if (!opt.empty())
+    exports.set_noise_out(opt[0][0]);
+  opt = get_options("rank_input");
+  if (!opt.empty())
+    exports.set_rank_input(opt[0][0]);
+  opt = get_options("rank_output");
   if (!opt.empty()) {
-    if (aggregator == aggregator_type::EXCLUSIVE) {
-      WARN("When using -aggregator exclusive, "
-           "the output of -weightedrank will be identical to the output of -rank, "
-           "as there is no aggregation of multiple patches per output voxel");
+    if (aggregator == aggregator_type::EXCLUSIVE && filter == filter_type::TRUNCATE) {
+      WARN("When using -aggregator exclusive and -filter truncate, "
+           "the output of -rank_output will be identical to the output of -rank_input, "
+           "as there is no aggregation of multiple patches per output voxel "
+           "and no optimal shrinkage to reduce output rank relative to estimated input rank");
     }
-    Header header(H3D);
-    header.datatype() = DataType::Float32;
-    header.datatype().set_byte_order_native();
-    weighted_rank = Image<float>::create(opt[0][0], header);
+    exports.set_rank_output(opt[0][0]);
   }
-
-  Image<float> sum_weights;
-  opt = get_options("sumweights");
+  opt = get_options("sum_optshrink");
   if (!opt.empty()) {
-    Header header(H3D);
-    header.datatype() = DataType::Float32;
-    header.datatype().set_byte_order_native();
-    sum_weights = Image<float>::create(opt[0][0], header);
     if (filter == filter_type::TRUNCATE) {
       WARN("Note that with a truncation filter, "
-           "output image from -sumweights option will be equivalent to rank");
+           "output image from -sumweights option will be equivalent to rank_input");
     }
+    exports.set_sum_optshrink(opt[0][0]);
   }
-
-  Image<float> max_dist;
   opt = get_options("max_dist");
-  if (!opt.empty()) {
-    Header header(H3D);
-    header.datatype() = DataType::Float32;
-    header.datatype().set_byte_order_native();
-    max_dist = Image<float>::create(opt[0][0], header);
-  }
-
-  Image<uint16_t> voxels;
-  opt = get_options("voxels");
-  if (!opt.empty()) {
-    Header header(H3D);
-    header.datatype() = DataType::UInt16;
-    header.datatype().set_byte_order_native();
-    voxels = Image<uint16_t>::create(opt[0][0], header);
-  }
+  if (!opt.empty())
+    exports.set_max_dist(opt[0][0]);
+  opt = get_options("voxelcount");
+  if (!opt.empty())
+    exports.set_voxelcount(opt[0][0]);
 
-  Image<float> aggregation_sum;
-  Header header_aggregation(H3D);
-  header_aggregation.datatype() = DataType::Float32;
-  header_aggregation.datatype().set_byte_order_native();
-  opt = get_options("aggregation_sum");
+  opt = get_options("sum_aggregation");
   if (!opt.empty()) {
     if (aggregator == aggregator_type::EXCLUSIVE) {
-      WARN("Output from -aggregation_sum will just contain 1 for every voxel processed: "
-           "no patch aggregation takes place when output series comex exclusively from central patch");
+      WARN("Output from -sum_aggregation will just contain 1 for every voxel processed: "
+           "no patch aggregation takes place when output series comes exclusively from central patch");
     }
-    aggregation_sum = Image<float>::create(opt[0][0], header_aggregation);
+    exports.set_sum_aggregation(opt[0][0]);
   } else if (aggregator != aggregator_type::EXCLUSIVE) {
-    aggregation_sum = Image<float>::scratch(header_aggregation, "Scratch buffer for patch aggregation weights");
-  }
-
-  opt = get_options("shape");
-  const shape_type shape = opt.empty() ? shape_type::SPHERE : shape_type((int)(opt[0][0]));
-  std::shared_ptr<KernelBase> kernel;
-
-  switch (shape) {
-  case shape_type::SPHERE: {
-    // TODO Could infer that user wants a cuboid kernel if -extent is used, even if -shape is not
-    if (!get_options("extent").empty())
-      throw Exception("-extent option does not apply to spherical kernel");
-    opt = get_options("radius_mm");
-    if (opt.empty())
-      kernel = std::make_shared<KernelSphereRatio>(dwi, get_option_value("radius_ratio", sphere_multiplier_default));
-    else
-      kernel = std::make_shared<KernelSphereFixedRadius>(dwi, opt[0][0]);
-  } break;
-  case shape_type::CUBOID: {
-    if (!get_options("radius_mm").empty() || !get_options("radius_ratio").empty())
-      throw Exception("-radius_* options are inapplicable if cuboid kernel shape is selected");
-    opt = get_options("extent");
-    std::vector<uint32_t> extent;
-    if (!opt.empty()) {
-      extent = parse_ints<uint32_t>(opt[0][0]);
-      if (extent.size() == 1)
-        extent = {extent[0], extent[0], extent[0]};
-      if (extent.size() != 3)
-        throw Exception("-extent must be either a scalar or a list of length 3");
-      for (int i = 0; i < 3; i++) {
-        if ((extent[i] & 1) == 0)
-          throw Exception("-extent must be a (list of) odd numbers");
-        if (extent[i] > dwi.size(i))
-          throw Exception("-extent must not exceed the image dimensions");
-      }
-    } else {
-      uint32_t e = 1;
-      while (Math::pow3(e) < dwi.size(3))
-        e += 2;
-      extent = {std::min(e, uint32_t(dwi.size(0))),  //
-                std::min(e, uint32_t(dwi.size(1))),  //
-                std::min(e, uint32_t(dwi.size(2)))}; //
-    }
-    INFO("selected patch size: " + str(extent[0]) + " x " + str(extent[1]) + " x " + str(extent[2]) + ".");
-
-    if (std::min<uint32_t>(dwi.size(3), extent[0] * extent[1] * extent[2]) < 15) {
-      WARN("The number of volumes or the patch size is small. "
-           "This may lead to discretisation effects in the noise level "
-           "and cause inconsistent denoising between adjacent voxels.");
-    }
-
-    kernel = std::make_shared<KernelCube>(dwi, extent);
-  } break;
-  default:
-    assert(false);
+    exports.set_sum_aggregation("");
   }
-  assert(kernel);
 
   int prec = get_option_value("datatype", 0); // default: single precision
   if (dwi.datatype().is_complex())
@@ -1053,71 +342,19 @@ void run() {
   switch (prec) {
   case 0:
     INFO("select real float32 for processing");
-    run<float>(dwi,
-               mask,
-               noise,
-               rank,
-               weighted_rank,
-               sum_weights,
-               max_dist,
-               voxels,
-               aggregation_sum,
-               argument[1],
-               kernel,
-               filter,
-               aggregator,
-               estimator);
+    run<float>(dwi, mask, kernel, estimator, filter, aggregator, argument[1], exports);
     break;
   case 1:
     INFO("select real float64 for processing");
-    run<double>(dwi,
-                mask,
-                noise,
-                rank,
-                weighted_rank,
-                sum_weights,
-                max_dist,
-                voxels,
-                aggregation_sum,
-                argument[1],
-                kernel,
-                filter,
-                aggregator,
-                estimator);
+    run<double>(dwi, mask, kernel, estimator, filter, aggregator, argument[1], exports);
     break;
   case 2:
     INFO("select complex float32 for processing");
-    run<cfloat>(dwi,
-                mask,
-                noise,
-                rank,
-                weighted_rank,
-                sum_weights,
-                max_dist,
-                voxels,
-                aggregation_sum,
-                argument[1],
-                kernel,
-                filter,
-                aggregator,
-                estimator);
+    run<cfloat>(dwi, mask, kernel, estimator, filter, aggregator, argument[1], exports);
     break;
   case 3:
     INFO("select complex float64 for processing");
-    run<cdouble>(dwi,
-                 mask,
-                 noise,
-                 rank,
-                 weighted_rank,
-                 sum_weights,
-                 max_dist,
-                 voxels,
-                 aggregation_sum,
-                 argument[1],
-                 kernel,
-                 filter,
-                 aggregator,
-                 estimator);
+    run<cdouble>(dwi, mask, kernel, estimator, filter, aggregator, argument[1], exports);
     break;
   }
 }
diff --git a/src/denoise/denoise.cpp b/src/denoise/denoise.cpp
new file mode 100644
index 0000000000..9f86a24267
--- /dev/null
+++ b/src/denoise/denoise.cpp
@@ -0,0 +1,30 @@
+/* Copyright (c) 2008-2024 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#include "denoise/denoise.h"
+
+namespace MR::Denoise {
+
+using namespace App;
+
+const Option datatype_option = Option("datatype",
+                                      "Datatype for the eigenvalue decomposition"
+                                      " (single or double precision). "
+                                      "For complex input data,"
+                                      " this will select complex float32 or complex float64 datatypes.") +
+                               Argument("float32/float64").type_choice(dtypes);
+
+} // namespace MR::Denoise
diff --git a/src/denoise/denoise.h b/src/denoise/denoise.h
new file mode 100644
index 0000000000..b35fa7f019
--- /dev/null
+++ b/src/denoise/denoise.h
@@ -0,0 +1,37 @@
+/* Copyright (c) 2008-2024 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#pragma once
+
+#include <Eigen/Dense>
+
+#include "app.h"
+
+namespace MR::Denoise {
+
+using eigenvalues_type = Eigen::Matrix<double, Eigen::Dynamic, 1>;
+using vector_type = Eigen::Array<double, Eigen::Dynamic, 1>;
+
+const std::vector<std::string> dtypes = {"float32", "float64"};
+extern const App::Option datatype_option;
+
+const std::vector<std::string> filters = {"truncate", "frobenius"};
+enum class filter_type { TRUNCATE, FROBENIUS };
+
+const std::vector<std::string> aggregators = {"exclusive", "gaussian", "invl0", "rank", "uniform"};
+enum class aggregator_type { EXCLUSIVE, GAUSSIAN, INVL0, RANK, UNIFORM };
+
+} // namespace MR::Denoise
diff --git a/src/denoise/estimator/base.h b/src/denoise/estimator/base.h
new file mode 100644
index 0000000000..8b239e1850
--- /dev/null
+++ b/src/denoise/estimator/base.h
@@ -0,0 +1,30 @@
+/* Copyright (c) 2008-2024 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#pragma once
+
+#include "denoise/denoise.h"
+#include "denoise/estimator/result.h"
+
+namespace MR::Denoise::Estimator {
+
+class Base {
+public:
+  Base() = default;
+  virtual Result operator()(const eigenvalues_type &eigenvalues, const ssize_t m, const ssize_t n) const = 0;
+};
+
+} // namespace MR::Denoise::Estimator
diff --git a/src/denoise/estimator/estimator.cpp b/src/denoise/estimator/estimator.cpp
new file mode 100644
index 0000000000..51efed00c7
--- /dev/null
+++ b/src/denoise/estimator/estimator.cpp
@@ -0,0 +1,54 @@
+/* Copyright (c) 2008-2024 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#include "denoise/estimator/estimator.h"
+
+#include "denoise/estimator/base.h"
+#include "denoise/estimator/exp.h"
+#include "denoise/estimator/mrm2022.h"
+
+namespace MR::Denoise::Estimator {
+
+using namespace App;
+
+const Option option = Option("estimator",
+                             "Select the noise level estimator"
+                             " (default = Exp2),"
+                             " either: \n"
+                             "* Exp1: the original estimator used in Veraart et al. (2016); \n"
+                             "* Exp2: the improved estimator introduced in Cordero-Grande et al. (2019); \n"
+                             "* MRM2022: the alternative estimator introduced in Olesen et al. (2022).") +
+                      Argument("algorithm").type_choice(estimators);
+
+std::shared_ptr<Base> make_estimator() {
+  auto opt = App::get_options("estimator");
+  const estimator_type est = opt.empty() ? estimator_type::EXP2 : estimator_type((int)(opt[0][0]));
+  switch (est) {
+  case estimator_type::EXP1:
+    return std::make_shared<Exp<1>>();
+  case estimator_type::EXP2:
+    return std::make_shared<Exp<2>>();
+    break;
+  case estimator_type::MRM2022:
+    return std::make_shared<MRM2022>();
+    break;
+  default:
+    assert(false);
+  }
+  return nullptr;
+}
+
+} // namespace MR::Denoise::Estimator
diff --git a/src/denoise/estimator/estimator.h b/src/denoise/estimator/estimator.h
new file mode 100644
index 0000000000..f0ec397949
--- /dev/null
+++ b/src/denoise/estimator/estimator.h
@@ -0,0 +1,34 @@
+/* Copyright (c) 2008-2024 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#pragma once
+
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "app.h"
+
+namespace MR::Denoise::Estimator {
+
+class Base;
+
+extern const App::Option option;
+const std::vector<std::string> estimators = {"exp1", "exp2", "mrm2022"};
+enum class estimator_type { EXP1, EXP2, MRM2022 };
+std::shared_ptr<Base> make_estimator();
+
+} // namespace MR::Denoise::Estimator
diff --git a/src/denoise/estimator/exp.h b/src/denoise/estimator/exp.h
new file mode 100644
index 0000000000..4931fe2b01
--- /dev/null
+++ b/src/denoise/estimator/exp.h
@@ -0,0 +1,62 @@
+/* Copyright (c) 2008-2024 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#pragma once
+
+#include "denoise/estimator/base.h"
+#include "denoise/estimator/result.h"
+
+namespace MR::Denoise::Estimator {
+
+// TODO Move to .cpp
+template <ssize_t version> class Exp : public Base {
+public:
+  Exp() = default;
+  Result operator()(const eigenvalues_type &s, const ssize_t m, const ssize_t n) const final {
+    Result result;
+    const ssize_t r = std::min(m, n);
+    const ssize_t q = std::max(m, n);
+    const double lam_r = std::max(s[0], 0.0) / q;
+    double clam = 0.0;
+    for (ssize_t p = 0; p < r; ++p) // p+1 is the number of noise components
+    {                               // (as opposed to the paper where p is defined as the number of signal components)
+      const double lam = std::max(s[p], 0.0) / q;
+      clam += lam;
+      double denominator = std::numeric_limits<double>::signaling_NaN();
+      switch (version) {
+      case 1:
+        denominator = q;
+        break;
+      case 2:
+        denominator = q - (r - p - 1);
+        break;
+      default:
+        assert(false);
+      }
+      const double gam = double(p + 1) / denominator;
+      const double sigsq1 = clam / double(p + 1);
+      const double sigsq2 = (lam - lam_r) / (4.0 * std::sqrt(gam));
+      // sigsq2 > sigsq1 if signal else noise
+      if (sigsq2 < sigsq1) {
+        result.sigma2 = sigsq1;
+        result.cutoff_p = p + 1;
+      }
+    }
+    return result;
+  }
+};
+
+} // namespace MR::Denoise::Estimator
diff --git a/src/denoise/estimator/mrm2022.h b/src/denoise/estimator/mrm2022.h
new file mode 100644
index 0000000000..1665ccb90c
--- /dev/null
+++ b/src/denoise/estimator/mrm2022.h
@@ -0,0 +1,58 @@
+/* Copyright (c) 2008-2024 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#pragma once
+
+#include <limits>
+
+#include "denoise/estimator/base.h"
+#include "denoise/estimator/result.h"
+#include "math/math.h"
+
+namespace MR::Denoise::Estimator {
+
+class MRM2022 : public Base {
+public:
+  MRM2022() = default;
+  Result operator()(const eigenvalues_type &s, const ssize_t m, const ssize_t n) const final {
+    Result result;
+    const ssize_t mprime = std::min(m, n);
+    const ssize_t nprime = std::max(m, n);
+    const double sigmasq_to_lamplus = Math::pow2(std::sqrt(nprime) + std::sqrt(mprime));
+    double clam = 0.0;
+    for (ssize_t i = 0; i != mprime; ++i)
+      clam += std::max(s[i], 0.0);
+    clam /= nprime;
+    // Unlike Exp# code,
+    //   MRM2022 article uses p to index number of signal components,
+    //   and here doing a direct translation of the manuscript content to code
+    double lamplusprev = -std::numeric_limits<double>::infinity();
+    for (ssize_t p = 0; p < mprime; ++p) {
+      const ssize_t i = mprime - 1 - p;
+      const double lam = std::max(s[i], 0.0) / nprime;
+      if (lam < lamplusprev)
+        return result;
+      clam -= lam;
+      const double sigmasq = clam / ((mprime - p) * (nprime - p));
+      lamplusprev = sigmasq * sigmasq_to_lamplus;
+      result.cutoff_p = i;
+      result.sigma2 = sigmasq;
+    }
+    return result;
+  }
+};
+
+} // namespace MR::Denoise::Estimator
diff --git a/src/denoise/estimator/result.h b/src/denoise/estimator/result.h
new file mode 100644
index 0000000000..94b511301f
--- /dev/null
+++ b/src/denoise/estimator/result.h
@@ -0,0 +1,28 @@
+/* Copyright (c) 2008-2024 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#pragma once
+
+namespace MR::Denoise::Estimator {
+
+class Result {
+public:
+  Result() : cutoff_p(0), sigma2(0.0) {}
+  ssize_t cutoff_p;
+  double sigma2;
+};
+
+} // namespace MR::Denoise::Estimator
diff --git a/src/denoise/exports.h b/src/denoise/exports.h
new file mode 100644
index 0000000000..72ee180e2a
--- /dev/null
+++ b/src/denoise/exports.h
@@ -0,0 +1,57 @@
+/* Copyright (c) 2008-2024 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#pragma once
+
+#include <string>
+
+#include "header.h"
+#include "image.h"
+
+namespace MR::Denoise {
+
+class Exports {
+public:
+  Exports(const Header &in) : H(in) {
+    H.ndim() = 3;
+    H.reset_intensity_scaling();
+  }
+  void set_noise_out(const std::string &path) { noise_out = Image<float>::create(path, H); }
+  void set_rank_input(const std::string &path) { rank_input = Image<uint16_t>::create(path, H); }
+  void set_rank_output(const std::string &path) { rank_output = Image<float>::create(path, H); }
+  void set_sum_optshrink(const std::string &path) { sum_optshrink = Image<float>::create(path, H); }
+  void set_max_dist(const std::string &path) { max_dist = Image<float>::create(path, H); }
+  void set_voxelcount(const std::string &path) { voxelcount = Image<uint16_t>::create(path, H); }
+  void set_sum_aggregation(const std::string &path) {
+    if (path.empty())
+      sum_aggregation = Image<float>::scratch(H, "Scratch image for patch aggregation sums");
+    else
+      sum_aggregation = Image<float>::create(path, H);
+  }
+
+  Image<float> noise_out;
+  Image<uint16_t> rank_input;
+  Image<float> rank_output;
+  Image<float> sum_optshrink;
+  Image<float> max_dist;
+  Image<uint16_t> voxelcount;
+  Image<float> sum_aggregation;
+
+protected:
+  Header H;
+};
+
+} // namespace MR::Denoise
diff --git a/src/denoise/functor.cpp b/src/denoise/functor.cpp
new file mode 100644
index 0000000000..4fab62d51f
--- /dev/null
+++ b/src/denoise/functor.cpp
@@ -0,0 +1,230 @@
+/* Copyright (c) 2008-2024 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#include "denoise/functor.h"
+
+#include "math/math.h"
+
+namespace MR::Denoise {
+
+template <typename F>
+Functor<F>::Functor(const Header &header,
+                    Image<bool> &mask,
+                    std::shared_ptr<Kernel::Base> kernel,
+                    std::shared_ptr<Estimator::Base> estimator,
+                    filter_type filter,
+                    aggregator_type aggregator,
+                    Exports &exports)
+    : m(header.size(3)),
+      mask(mask),
+      kernel(kernel),
+      estimator(estimator),
+      filter(filter),
+      aggregator(aggregator),
+      // FWHM = 2 x cube root of voxel spacings
+      gaussian_multiplier(-std::log(2.0) /
+                          Math::pow2(std::cbrt(header.spacing(0) * header.spacing(1) * header.spacing(2)))),
+      X(m, kernel->estimated_size()),
+      XtX(std::min(m, kernel->estimated_size()), std::min(m, kernel->estimated_size())),
+      eig(std::min(m, kernel->estimated_size())),
+      s(std::min(m, kernel->estimated_size())),
+      clam(std::min(m, kernel->estimated_size())),
+      w(std::min(m, kernel->estimated_size())),
+      exports(exports) {}
+
+template <typename F> void Functor<F>::operator()(Image<F> &dwi, Image<F> &out) {
+  // Process voxels in mask only
+  if (mask.valid()) {
+    assign_pos_of(dwi, 0, 3).to(mask);
+    if (!mask.value())
+      return;
+  }
+
+  // Load list of voxels from which to load data
+  const Kernel::Data neighbourhood = (*kernel)({dwi.index(0), dwi.index(1), dwi.index(2)});
+  const ssize_t n = neighbourhood.voxels.size();
+  const ssize_t r = std::min(m, n);
+  const ssize_t q = std::max(m, n);
+
+  // Expand local storage if necessary
+  if (n > X.cols()) {
+    DEBUG("Expanding data matrix storage from " + str(m) + "x" + str(X.cols()) + " to " + str(m) + "x" + str(n));
+    X.resize(m, n);
+  }
+  if (r > XtX.cols()) {
+    DEBUG("Expanding decomposition matrix storage from " + str(X.rows()) + " to " + str(r));
+    XtX.resize(r, r);
+    s.resize(r);
+    clam.resize(r);
+    w.resize(r);
+  }
+
+  // Fill matrices with NaN when in debug mode;
+  //   make sure results from one voxel are not creeping into another
+  //   due to use of block oberations to prevent memory re-allocation
+  //   in the presence of variation in kernel sizes
+#ifndef NDEBUG
+  X.fill(std::numeric_limits<F>::signaling_NaN());
+  XtX.fill(std::numeric_limits<F>::signaling_NaN());
+  s.fill(std::numeric_limits<default_type>::signaling_NaN());
+  clam.fill(std::numeric_limits<default_type>::signaling_NaN());
+  w.fill(std::numeric_limits<default_type>::signaling_NaN());
+#endif
+
+  load_data(dwi, neighbourhood.voxels);
+
+  // Compute Eigendecomposition:
+  if (m <= n)
+    XtX.topLeftCorner(r, r).template triangularView<Eigen::Lower>() = X.leftCols(n) * X.leftCols(n).adjoint();
+  else
+    XtX.topLeftCorner(r, r).template triangularView<Eigen::Lower>() = X.leftCols(n).adjoint() * X.leftCols(n);
+  eig.compute(XtX.topLeftCorner(r, r));
+  // eigenvalues sorted in increasing order:
+  s.head(r) = eig.eigenvalues().template cast<double>();
+
+  // Marchenko-Pastur optimal threshold determination
+  const Estimator::Result threshold = (*estimator)(s, m, n);
+
+  // Generate weights vector
+  double sum_weights = 0.0;
+  ssize_t out_rank = 0;
+  switch (filter) {
+  case filter_type::TRUNCATE:
+    out_rank = r - threshold.cutoff_p;
+    w.head(threshold.cutoff_p).setZero();
+    w.segment(threshold.cutoff_p, r - threshold.cutoff_p).setOnes();
+    sum_weights = double(out_rank);
+    break;
+  case filter_type::FROBENIUS: {
+    const double beta = r / q;
+    const double transition = 1.0 + std::sqrt(beta);
+    double clam = 0.0;
+    for (ssize_t i = 0; i != r; ++i) {
+      const double lam = std::max(s[i], 0.0) / q;
+      clam += lam;
+      const double y = clam / (threshold.sigma2 * (i + 1));
+      double nu = 0.0;
+      if (y > transition) {
+        nu = std::sqrt(Math::pow2(Math::pow2(y) - beta - 1.0) - (4.0 * beta)) / y;
+        ++out_rank;
+      }
+      w[i] = nu / y;
+      sum_weights += w[i];
+    }
+  } break;
+  default:
+    assert(false);
+  }
+
+  // recombine data using only eigenvectors above threshold
+  // If only the data computed when this voxel was the centre of the patch
+  //   is to be used for synthesis of the output image,
+  //   then only that individual column needs to be reconstructed;
+  //   if however the result from this patch is to contribute to the synthesized image
+  //   for all voxels that were utilised within this patch,
+  //   then we need to instead compute the full projection
+  switch (aggregator) {
+  case aggregator_type::EXCLUSIVE:
+    if (m <= n)
+      X.col(neighbourhood.centre_index) =
+          eig.eigenvectors() * (w.head(r).cast<F>().matrix().asDiagonal() *
+                                (eig.eigenvectors().adjoint() * X.col(neighbourhood.centre_index)));
+    else
+      X.col(neighbourhood.centre_index) =
+          X.leftCols(n) * (eig.eigenvectors() * (w.head(r).cast<F>().matrix().asDiagonal() *
+                                                 eig.eigenvectors().adjoint().col(neighbourhood.centre_index)));
+    assign_pos_of(dwi).to(out);
+    out.row(3) = X.col(neighbourhood.centre_index);
+    if (exports.sum_aggregation.valid()) {
+      assign_pos_of(dwi, 0, 3).to(exports.sum_aggregation);
+      exports.sum_aggregation.value() = 1.0;
+    }
+    if (exports.rank_output.valid()) {
+      assign_pos_of(dwi, 0, 3).to(exports.rank_output);
+      exports.rank_output.value() = out_rank;
+    }
+    break;
+  default: {
+    if (m <= n)
+      X = eig.eigenvectors() * (w.head(r).cast<F>().matrix().asDiagonal() * (eig.eigenvectors().adjoint() * X));
+    else
+      X.leftCols(n) = X.leftCols(n) *
+                      (eig.eigenvectors() * (w.head(r).cast<F>().matrix().asDiagonal() * eig.eigenvectors().adjoint()));
+    std::lock_guard<std::mutex> lock(mutex_aggregator);
+    for (size_t voxel_index = 0; voxel_index != neighbourhood.voxels.size(); ++voxel_index) {
+      assign_pos_of(neighbourhood.voxels[voxel_index].index, 0, 3).to(out);
+      assign_pos_of(neighbourhood.voxels[voxel_index].index).to(exports.sum_aggregation);
+      double weight = std::numeric_limits<double>::signaling_NaN();
+      switch (aggregator) {
+      case aggregator_type::EXCLUSIVE:
+        assert(false);
+        break;
+      case aggregator_type::GAUSSIAN:
+        weight = std::exp(gaussian_multiplier * neighbourhood.voxels[voxel_index].sq_distance);
+        break;
+      case aggregator_type::INVL0:
+        weight = 1.0 / (1 + out_rank);
+        break;
+      case aggregator_type::RANK:
+        weight = out_rank;
+        break;
+      case aggregator_type::UNIFORM:
+        weight = 1.0;
+        break;
+      }
+      out.row(3) += weight * X.col(voxel_index);
+      exports.sum_aggregation.value() += weight;
+      if (exports.rank_output.valid()) {
+        assign_pos_of(neighbourhood.voxels[voxel_index].index, 0, 3).to(exports.rank_output);
+        exports.rank_output.value() += weight * out_rank;
+      }
+    }
+  } break;
+  }
+
+  // Store additional output maps if requested
+  if (exports.noise_out.valid()) {
+    assign_pos_of(dwi, 0, 3).to(exports.noise_out);
+    exports.noise_out.value() = float(std::sqrt(threshold.sigma2));
+  }
+  if (exports.rank_input.valid()) {
+    assign_pos_of(dwi, 0, 3).to(exports.rank_input);
+    exports.rank_input.value() = out_rank;
+  }
+  if (exports.sum_optshrink.valid()) {
+    assign_pos_of(dwi, 0, 3).to(exports.sum_optshrink);
+    exports.sum_optshrink.value() = sum_weights;
+  }
+  if (exports.max_dist.valid()) {
+    assign_pos_of(dwi, 0, 3).to(exports.max_dist);
+    exports.max_dist.value() = neighbourhood.max_distance;
+  }
+  if (exports.voxelcount.valid()) {
+    assign_pos_of(dwi, 0, 3).to(exports.voxelcount);
+    exports.voxelcount.value() = n;
+  }
+}
+
+template <typename F> void Functor<F>::load_data(Image<F> &image, const std::vector<Kernel::Voxel> &voxels) {
+  const Kernel::Voxel::index_type pos({image.index(0), image.index(1), image.index(2)});
+  for (ssize_t i = 0; i != voxels.size(); ++i) {
+    assign_pos_of(voxels[i].index, 0, 3).to(image);
+    X.col(i) = image.row(3);
+  }
+  assign_pos_of(pos, 0, 3).to(image);
+}
+
+} // namespace MR::Denoise
diff --git a/src/denoise/functor.h b/src/denoise/functor.h
new file mode 100644
index 0000000000..11cf94c42a
--- /dev/null
+++ b/src/denoise/functor.h
@@ -0,0 +1,91 @@
+/* Copyright (c) 2008-2024 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#pragma once
+
+#include <limits>
+#include <memory>
+#include <mutex>
+
+#include <Eigen/Dense>
+
+#include "denoise/denoise.h"
+#include "denoise/estimator/base.h"
+#include "denoise/estimator/result.h"
+#include "denoise/exports.h"
+#include "denoise/kernel/base.h"
+#include "denoise/kernel/data.h"
+#include "denoise/kernel/voxel.h"
+#include "header.h"
+#include "image.h"
+
+namespace MR::Denoise {
+
+template <typename F> class Functor {
+
+public:
+  using MatrixType = Eigen::Matrix<F, Eigen::Dynamic, Eigen::Dynamic>;
+
+  Functor(const Header &header,
+          Image<bool> &mask,
+          std::shared_ptr<Kernel::Base> kernel,
+          std::shared_ptr<Estimator::Base> estimator,
+          filter_type filter,
+          aggregator_type aggregator,
+          Exports &exports);
+
+  void operator()(Image<F> &dwi, Image<F> &out);
+
+private:
+  // Denoising configuration
+  const ssize_t m;
+  Image<bool> mask;
+  std::shared_ptr<Kernel::Base> kernel;
+  std::shared_ptr<Estimator::Base> estimator;
+  filter_type filter;
+  aggregator_type aggregator;
+  double gaussian_multiplier;
+
+  // Reusable memory
+  MatrixType X;
+  MatrixType XtX;
+  Eigen::SelfAdjointEigenSolver<MatrixType> eig;
+  eigenvalues_type s;
+  vector_type clam;
+  vector_type w;
+
+  // Export images
+  Exports exports;
+
+  // Data that can only be written in a thread-safe manner
+  // Note that this applies not just to this scratch buffer, but also the output image
+  //   (while it would be thread-safe to create a full copy of the output image for each thread
+  //   and combine them only at destruction time,
+  //   this runs the risk of becoming prohibitively large)
+  // Not placing this within a MutexProtexted<> as the image type is still templated
+  static std::mutex mutex_aggregator;
+
+  void load_data(Image<F> &image, const std::vector<Kernel::Voxel> &voxels);
+};
+
+template <typename F> std::mutex Functor<F>::mutex_aggregator;
+
+template class Functor<float>;
+template class Functor<cfloat>;
+template class Functor<double>;
+template class Functor<cdouble>;
+
+} // namespace MR::Denoise
diff --git a/src/denoise/kernel/base.h b/src/denoise/kernel/base.h
new file mode 100644
index 0000000000..0e8d47b7ed
--- /dev/null
+++ b/src/denoise/kernel/base.h
@@ -0,0 +1,39 @@
+/* Copyright (c) 2008-2024 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#pragma once
+
+#include "denoise/kernel/data.h"
+#include "denoise/kernel/voxel.h"
+#include "header.h"
+
+namespace MR::Denoise::Kernel {
+
+class Base {
+public:
+  Base(const Header &H) : H(H) {}
+  Base(const Base &) = default;
+  virtual ~Base() = default;
+  // This is just for pre-allocating matrices
+  virtual ssize_t estimated_size() const = 0;
+  // This is the interface that kernels must provide
+  virtual Data operator()(const Voxel::index_type &) const = 0;
+
+protected:
+  const Header H;
+};
+
+} // namespace MR::Denoise::Kernel
diff --git a/src/denoise/kernel/cuboid.cpp b/src/denoise/kernel/cuboid.cpp
new file mode 100644
index 0000000000..261c97fbfa
--- /dev/null
+++ b/src/denoise/kernel/cuboid.cpp
@@ -0,0 +1,67 @@
+/* Copyright (c) 2008-2024 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#include "denoise/kernel/cuboid.h"
+
+namespace MR::Denoise::Kernel {
+
+Cuboid::Cuboid(const Header &header, const std::vector<uint32_t> &extent)
+    : Base(header),
+      half_extent({ssize_t(extent[0] / 2), ssize_t(extent[1] / 2), ssize_t(extent[2] / 2)}),
+      size(ssize_t(extent[0]) * ssize_t(extent[1]) * ssize_t(extent[2])),
+      centre_index(size / 2) {
+  for (auto e : extent) {
+    if (!(e % 2))
+      throw Exception("Size of cubic kernel must be an odd integer");
+  }
+}
+
+namespace {
+// patch handling at image edges
+inline ssize_t wrapindex(int p, int r, int e, int max) {
+  int rr = p + r;
+  if (rr < 0)
+    rr = e - r;
+  if (rr >= max)
+    rr = (max - 1) - e - r;
+  return rr;
+}
+} // namespace
+
+Data Cuboid::operator()(const Voxel::index_type &pos) const {
+  Data result(centre_index);
+  Voxel::index_type voxel;
+  Offset::index_type offset;
+  for (offset[2] = -half_extent[2]; offset[2] <= half_extent[2]; ++offset[2]) {
+    voxel[2] = wrapindex(pos[2], offset[2], half_extent[2], H.size(2));
+    for (offset[1] = -half_extent[1]; offset[1] <= half_extent[1]; ++offset[1]) {
+      voxel[1] = wrapindex(pos[1], offset[1], half_extent[1], H.size(1));
+      for (offset[0] = -half_extent[0]; offset[0] <= half_extent[0]; ++offset[0]) {
+        voxel[0] = wrapindex(pos[0], offset[0], half_extent[0], H.size(0));
+        // Both "pos" and "voxel" are unsigned, so beware of integer overflow
+        const default_type sq_distance = Math::pow2(std::min(pos[0] - voxel[0], voxel[0] - pos[0]) * H.spacing(0)) +
+                                         Math::pow2(std::min(pos[1] - voxel[1], voxel[1] - pos[1]) * H.spacing(1)) +
+                                         Math::pow2(std::min(pos[2] - voxel[2], voxel[2] - pos[2]) * H.spacing(2));
+        result.voxels.push_back(Voxel(voxel, sq_distance));
+        result.max_distance = std::max(result.max_distance, sq_distance);
+      }
+    }
+  }
+  result.max_distance = std::sqrt(result.max_distance);
+  return result;
+}
+
+} // namespace MR::Denoise::Kernel
diff --git a/src/denoise/kernel/cuboid.h b/src/denoise/kernel/cuboid.h
new file mode 100644
index 0000000000..0d2230009d
--- /dev/null
+++ b/src/denoise/kernel/cuboid.h
@@ -0,0 +1,42 @@
+/* Copyright (c) 2008-2024 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#pragma once
+
+#include <vector>
+
+#include "denoise/kernel/base.h"
+#include "denoise/kernel/data.h"
+#include "header.h"
+
+namespace MR::Denoise::Kernel {
+
+class Cuboid : public Base {
+
+public:
+  Cuboid(const Header &header, const std::vector<uint32_t> &extent);
+  Cuboid(const Cuboid &) = default;
+  ~Cuboid() final = default;
+  Data operator()(const Voxel::index_type &pos) const override;
+  ssize_t estimated_size() const override { return size; }
+
+private:
+  const Voxel::index_type half_extent;
+  const ssize_t size;
+  const ssize_t centre_index;
+};
+
+} // namespace MR::Denoise::Kernel
diff --git a/src/denoise/kernel/data.h b/src/denoise/kernel/data.h
new file mode 100644
index 0000000000..5d783baa05
--- /dev/null
+++ b/src/denoise/kernel/data.h
@@ -0,0 +1,35 @@
+/* Copyright (c) 2008-2024 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#pragma once
+
+#include <vector>
+
+#include "denoise/kernel/voxel.h"
+#include "types.h"
+
+namespace MR::Denoise::Kernel {
+
+class Data {
+public:
+  Data() : centre_index(-1), max_distance(-std::numeric_limits<default_type>::infinity()) {}
+  Data(const ssize_t i) : centre_index(i), max_distance(-std::numeric_limits<default_type>::infinity()) {}
+  std::vector<Voxel> voxels;
+  ssize_t centre_index;
+  default_type max_distance;
+};
+
+} // namespace MR::Denoise::Kernel
diff --git a/src/denoise/kernel/kernel.cpp b/src/denoise/kernel/kernel.cpp
new file mode 100644
index 0000000000..e5cd077166
--- /dev/null
+++ b/src/denoise/kernel/kernel.cpp
@@ -0,0 +1,128 @@
+/* Copyright (c) 2008-2024 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#include "denoise/kernel/kernel.h"
+
+#include "denoise/kernel/base.h"
+#include "denoise/kernel/cuboid.h"
+#include "denoise/kernel/sphere_radius.h"
+#include "denoise/kernel/sphere_ratio.h"
+#include "math/math.h"
+
+namespace MR::Denoise::Kernel {
+
+using namespace App;
+
+const char *const shape_description =
+    "The sliding spatial window behaves differently at the edges of the image FoV "
+    "depending on the shape / size selected for that window. "
+    "The default behaviour is to use a spherical kernel centred at the voxel of interest, "
+    "whose size is some multiple of the number of input volumes; "
+    "where some such voxels lie outside of the image FoV, "
+    "the radius of the kernel will be increased until the requisite number of voxels are used. "
+    "For a spherical kernel of a fixed radius, "
+    "no such expansion will occur, "
+    "and so for voxels near the image edge a reduced number of voxels will be present in the kernel. "
+    "For a cuboid kernel, "
+    "the centre of the kernel will be offset from the voxel being processed "
+    "such that the entire volume of the kernel resides within the image FoV.";
+
+const char *const size_description =
+    "The size of the default spherical kernel is set to select a number of voxels that is "
+    "1.0 / 0.85 ~ 1.18 times the number of volumes in the input series. "
+    "If a cuboid kernel is requested, "
+    "but the -extent option is not specified, "
+    "the command will select the smallest isotropic patch size "
+    "that exceeds the number of DW images in the input data; "
+    "e.g., 5x5x5 for data with <= 125 DWI volumes, "
+    "7x7x7 for data with <= 343 DWI volumes, etc.";
+
+// clang-format off
+const OptionGroup options = OptionGroup("Options for controlling the sliding spatial window kernel")
++ Option("shape",
+         "Set the shape of the sliding spatial window. "
+         "Options are: " + join(shapes, ",") + "; default: sphere")
+  + Argument("choice").type_choice(shapes)
++ Option("radius_mm", "Set an absolute spherical kernel radius in mm")
+  + Argument("value").type_float(0.0)
++ Option("radius_ratio",
+         "Set the spherical kernel size as a ratio of number of voxels to number of input volumes "
+         "(default: 1.0/0.85 ~= 1.18)")
+  + Argument("value").type_float(0.0)
+// TODO Command-line option that allows user to specify minimum absolute number of voxels in kernel
++ Option("extent",
+         "Set the patch size of the cuboid kernel; "
+         "can be either a single odd integer or a comma-separated triplet of odd integers")
+  + Argument("window").type_sequence_int();
+// clang-format on
+
+std::shared_ptr<Base> make_kernel(const Header &H) {
+  auto opt = App::get_options("shape");
+  const Kernel::shape_type shape = opt.empty() ? Kernel::shape_type::SPHERE : Kernel::shape_type((int)(opt[0][0]));
+  std::shared_ptr<Kernel::Base> kernel;
+
+  switch (shape) {
+  case Kernel::shape_type::SPHERE: {
+    // TODO Could infer that user wants a cuboid kernel if -extent is used, even if -shape is not
+    if (!get_options("extent").empty())
+      throw Exception("-extent option does not apply to spherical kernel");
+    opt = get_options("radius_mm");
+    if (opt.empty())
+      return std::make_shared<SphereRatio>(H, get_option_value("radius_ratio", sphere_multiplier_default));
+    return std::make_shared<SphereFixedRadius>(H, opt[0][0]);
+  }
+  case Kernel::shape_type::CUBOID: {
+    if (!get_options("radius_mm").empty() || !get_options("radius_ratio").empty())
+      throw Exception("-radius_* options are inapplicable if cuboid kernel shape is selected");
+    opt = get_options("extent");
+    std::vector<uint32_t> extent;
+    if (!opt.empty()) {
+      extent = parse_ints<uint32_t>(opt[0][0]);
+      if (extent.size() == 1)
+        extent = {extent[0], extent[0], extent[0]};
+      if (extent.size() != 3)
+        throw Exception("-extent must be either a scalar or a list of length 3");
+      for (int i = 0; i < 3; i++) {
+        if ((extent[i] & 1) == 0)
+          throw Exception("-extent must be a (list of) odd numbers");
+        if (extent[i] > H.size(i))
+          throw Exception("-extent must not exceed the image dimensions");
+      }
+    } else {
+      uint32_t e = 1;
+      while (Math::pow3(e) < H.size(3))
+        e += 2;
+      extent = {std::min(e, uint32_t(H.size(0))),  //
+                std::min(e, uint32_t(H.size(1))),  //
+                std::min(e, uint32_t(H.size(2)))}; //
+    }
+    INFO("selected patch size: " + str(extent[0]) + " x " + str(extent[1]) + " x " + str(extent[2]) + ".");
+
+    if (std::min<uint32_t>(H.size(3), extent[0] * extent[1] * extent[2]) < 15) {
+      WARN("The number of volumes or the patch size is small. "
+           "This may lead to discretisation effects in the noise level "
+           "and cause inconsistent denoising between adjacent voxels.");
+    }
+
+    return std::make_shared<Cuboid>(H, extent);
+  } break;
+  default:
+    assert(false);
+  }
+  return nullptr;
+}
+
+} // namespace MR::Denoise::Kernel
diff --git a/src/denoise/kernel/kernel.h b/src/denoise/kernel/kernel.h
new file mode 100644
index 0000000000..3bcfb3d2d3
--- /dev/null
+++ b/src/denoise/kernel/kernel.h
@@ -0,0 +1,39 @@
+/* Copyright (c) 2008-2024 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#pragma once
+
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "app.h"
+#include "header.h"
+#include "types.h"
+
+namespace MR::Denoise::Kernel {
+
+class Base;
+
+extern const char *const shape_description;
+extern const char *const size_description;
+
+const std::vector<std::string> shapes = {"cuboid", "sphere"};
+enum class shape_type { CUBOID, SPHERE };
+extern const App::OptionGroup options;
+std::shared_ptr<Base> make_kernel(const Header &H);
+
+} // namespace MR::Denoise::Kernel
diff --git a/src/denoise/kernel/sphere_base.cpp b/src/denoise/kernel/sphere_base.cpp
new file mode 100644
index 0000000000..0fadb90354
--- /dev/null
+++ b/src/denoise/kernel/sphere_base.cpp
@@ -0,0 +1,45 @@
+/* Copyright (c) 2008-2024 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#include "denoise/kernel/sphere_base.h"
+
+#include "math/math.h"
+
+namespace MR::Denoise::Kernel {
+
+SphereBase::Shared::Shared(const Header &voxel_grid, const default_type max_radius) {
+  const default_type max_radius_sq = Math::pow2(max_radius);
+  const Voxel::index_type half_extents({ssize_t(std::ceil(max_radius / voxel_grid.spacing(0))),   //
+                                        ssize_t(std::ceil(max_radius / voxel_grid.spacing(1))),   //
+                                        ssize_t(std::ceil(max_radius / voxel_grid.spacing(2)))}); //
+  // Build the searchlight
+  data.reserve(size_t(2 * half_extents[0] + 1) * size_t(2 * half_extents[1] + 1) * size_t(2 * half_extents[2] + 1));
+  Offset::index_type offset({0, 0, 0});
+  for (offset[2] = -half_extents[2]; offset[2] <= half_extents[2]; ++offset[2]) {
+    for (offset[1] = -half_extents[1]; offset[1] <= half_extents[1]; ++offset[1]) {
+      for (offset[0] = -half_extents[0]; offset[0] <= half_extents[0]; ++offset[0]) {
+        const default_type squared_distance = Math::pow2(offset[0] * voxel_grid.spacing(0))    //
+                                              + Math::pow2(offset[1] * voxel_grid.spacing(1))  //
+                                              + Math::pow2(offset[2] * voxel_grid.spacing(2)); //
+        if (squared_distance <= max_radius_sq)
+          data.emplace_back(Offset(offset, squared_distance));
+      }
+    }
+  }
+  std::sort(data.begin(), data.end());
+}
+
+} // namespace MR::Denoise::Kernel
diff --git a/src/denoise/kernel/sphere_base.h b/src/denoise/kernel/sphere_base.h
new file mode 100644
index 0000000000..7dbe7af942
--- /dev/null
+++ b/src/denoise/kernel/sphere_base.h
@@ -0,0 +1,54 @@
+/* Copyright (c) 2008-2024 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#pragma once
+
+#include <memory>
+#include <vector>
+
+#include "denoise/kernel/base.h"
+#include "denoise/kernel/kernel.h"
+#include "denoise/kernel/voxel.h"
+#include "header.h"
+
+namespace MR::Denoise::Kernel {
+
+class SphereBase : public Base {
+
+public:
+  SphereBase(const Header &voxel_grid, const default_type max_radius)
+      : Base(voxel_grid), shared(new Shared(voxel_grid, max_radius)) {}
+
+  SphereBase(const SphereBase &) = default;
+
+  virtual ~SphereBase() override {}
+
+protected:
+  class Shared {
+  public:
+    using TableType = std::vector<Offset>;
+    Shared(const Header &voxel_grid, const default_type max_radius);
+    TableType::const_iterator begin() const { return data.begin(); }
+    TableType::const_iterator end() const { return data.end(); }
+
+  private:
+    TableType data;
+  };
+
+  std::shared_ptr<Shared> shared;
+};
+
+} // namespace MR::Denoise::Kernel
diff --git a/src/denoise/kernel/sphere_radius.cpp b/src/denoise/kernel/sphere_radius.cpp
new file mode 100644
index 0000000000..d759ce715b
--- /dev/null
+++ b/src/denoise/kernel/sphere_radius.cpp
@@ -0,0 +1,37 @@
+/* Copyright (c) 2008-2024 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#include "denoise/kernel/sphere_radius.h"
+
+namespace MR::Denoise::Kernel {
+
+Data SphereFixedRadius::operator()(const Voxel::index_type &pos) const {
+  Data result(0);
+  result.voxels.reserve(maximum_size);
+  for (auto map_it = shared->begin(); map_it != shared->end(); ++map_it) {
+    const Voxel::index_type voxel({pos[0] + map_it->index[0],   //
+                                   pos[1] + map_it->index[1],   //
+                                   pos[2] + map_it->index[2]}); //
+    if (!is_out_of_bounds(H, voxel, 0, 3)) {
+      result.voxels.push_back(Voxel(voxel, map_it->sq_distance));
+      result.max_distance = map_it->sq_distance;
+    }
+  }
+  result.max_distance = std::sqrt(result.max_distance);
+  return result;
+}
+
+} // namespace MR::Denoise::Kernel
diff --git a/src/denoise/kernel/sphere_radius.h b/src/denoise/kernel/sphere_radius.h
new file mode 100644
index 0000000000..002dd8243a
--- /dev/null
+++ b/src/denoise/kernel/sphere_radius.h
@@ -0,0 +1,42 @@
+/* Copyright (c) 2008-2024 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#pragma once
+
+#include "denoise/kernel/data.h"
+#include "denoise/kernel/sphere_base.h"
+#include "header.h"
+#include "types.h"
+
+namespace MR::Denoise::Kernel {
+
+class SphereFixedRadius : public SphereBase {
+public:
+  SphereFixedRadius(const Header &voxel_grid, const default_type radius)
+      : SphereBase(voxel_grid, radius),                               //
+        maximum_size(std::distance(shared->begin(), shared->end())) { //
+    INFO("Maximum number of voxels in " + str(radius) + "mm fixed-radius kernel is " + str(maximum_size));
+  }
+  SphereFixedRadius(const SphereFixedRadius &) = default;
+  ~SphereFixedRadius() final = default;
+  Data operator()(const Voxel::index_type &pos) const;
+  ssize_t estimated_size() const override { return maximum_size; }
+
+private:
+  const ssize_t maximum_size;
+};
+
+} // namespace MR::Denoise::Kernel
diff --git a/src/denoise/kernel/sphere_ratio.cpp b/src/denoise/kernel/sphere_ratio.cpp
new file mode 100644
index 0000000000..f682c37198
--- /dev/null
+++ b/src/denoise/kernel/sphere_ratio.cpp
@@ -0,0 +1,62 @@
+/* Copyright (c) 2008-2024 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#include "denoise/kernel/sphere_ratio.h"
+
+namespace MR::Denoise::Kernel {
+
+Data SphereRatio::operator()(const Voxel::index_type &pos) const {
+  Data result(0);
+  auto table_it = shared->begin();
+  while (table_it != shared->end()) {
+    // If there's a tie in distances, want to include all such offsets in the kernel,
+    //   even if the size of the utilised kernel extends beyond the minimum size
+    if (result.voxels.size() >= min_size && table_it->sq_distance != result.max_distance)
+      break;
+    const Voxel::index_type voxel({pos[0] + table_it->index[0],   //
+                                   pos[1] + table_it->index[1],   //
+                                   pos[2] + table_it->index[2]}); //
+    if (!is_out_of_bounds(H, voxel, 0, 3)) {
+      result.voxels.push_back(Voxel(voxel, table_it->sq_distance));
+      result.max_distance = table_it->sq_distance;
+    }
+    ++table_it;
+  }
+  if (table_it == shared->end()) {
+    throw Exception(                                                                   //
+        std::string("Inadequate spherical kernel initialisation ")                     //
+        + "(lookup table " + str(std::distance(shared->begin(), shared->end())) + "; " //
+        + "min size " + str(min_size) + "; "                                           //
+        + "read size " + str(result.voxels.size()) + ")");                             //
+  }
+  result.max_distance = std::sqrt(result.max_distance);
+  return result;
+}
+
+default_type SphereRatio::compute_max_radius(const Header &voxel_grid, const default_type min_ratio) const {
+  const size_t num_volumes = voxel_grid.size(3);
+  const default_type voxel_volume = voxel_grid.spacing(0) * voxel_grid.spacing(1) * voxel_grid.spacing(2);
+  const default_type sphere_volume = 8.0 * num_volumes * min_ratio * voxel_volume;
+  const default_type approx_radius = std::sqrt(sphere_volume * 0.75 / Math::pi);
+  const Voxel::index_type half_extents({ssize_t(std::ceil(approx_radius / voxel_grid.spacing(0))),   //
+                                        ssize_t(std::ceil(approx_radius / voxel_grid.spacing(1))),   //
+                                        ssize_t(std::ceil(approx_radius / voxel_grid.spacing(2)))}); //
+  return std::max({half_extents[0] * voxel_grid.spacing(0),
+                   half_extents[1] * voxel_grid.spacing(1),
+                   half_extents[2] * voxel_grid.spacing(2)});
+}
+
+} // namespace MR::Denoise::Kernel
diff --git a/src/denoise/kernel/sphere_ratio.h b/src/denoise/kernel/sphere_ratio.h
new file mode 100644
index 0000000000..bb75268ed6
--- /dev/null
+++ b/src/denoise/kernel/sphere_ratio.h
@@ -0,0 +1,51 @@
+/* Copyright (c) 2008-2024 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#pragma once
+
+#include "denoise/kernel/data.h"
+#include "denoise/kernel/sphere_base.h"
+#include "header.h"
+
+namespace MR::Denoise::Kernel {
+
+constexpr default_type sphere_multiplier_default = 1.0 / 0.85;
+
+class SphereRatio : public SphereBase {
+
+public:
+  SphereRatio(const Header &voxel_grid, const default_type min_ratio)
+      : SphereBase(voxel_grid, compute_max_radius(voxel_grid, min_ratio)),
+        min_size(std::ceil(voxel_grid.size(3) * min_ratio)) {}
+
+  SphereRatio(const SphereRatio &) = default;
+
+  ~SphereRatio() final = default;
+
+  Data operator()(const Voxel::index_type &pos) const override;
+
+  ssize_t estimated_size() const override { return min_size; }
+
+private:
+  ssize_t min_size;
+
+  // Determine an appropriate bounding box from which to generate the search table
+  // Find the radius for which 7/8 of the sphere will contain the minimum number of voxels, then round up
+  // This is only for setting the maximal radius for generation of the lookup table
+  default_type compute_max_radius(const Header &voxel_grid, const default_type min_ratio) const;
+};
+
+} // namespace MR::Denoise::Kernel
diff --git a/src/denoise/kernel/voxel.h b/src/denoise/kernel/voxel.h
new file mode 100644
index 0000000000..c76ff286ce
--- /dev/null
+++ b/src/denoise/kernel/voxel.h
@@ -0,0 +1,54 @@
+/* Copyright (c) 2008-2024 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#pragma once
+
+#include <Eigen/Dense>
+
+#include "types.h"
+
+namespace MR::Denoise::Kernel {
+
+template <class T> class VoxelBase {
+public:
+  using index_type = Eigen::Array<T, 3, 1>;
+  VoxelBase(const index_type &index, const default_type sq_distance) : index(index), sq_distance(sq_distance) {}
+  VoxelBase(const VoxelBase &) = default;
+  VoxelBase(VoxelBase &&) = default;
+  ~VoxelBase() {}
+  VoxelBase &operator=(const VoxelBase &that) {
+    index = that.index;
+    sq_distance = that.sq_distance;
+    return *this;
+  }
+  VoxelBase &operator=(VoxelBase &&that) noexcept {
+    index = that.index;
+    sq_distance = that.sq_distance;
+    return *this;
+  }
+  bool operator<(const VoxelBase &that) const { return sq_distance < that.sq_distance; }
+  default_type distance() const { return std::sqrt(sq_distance); }
+
+  index_type index;
+  default_type sq_distance;
+};
+
+// Need signed integer to represent offsets from the centre of the kernel;
+//   however absolute voxel indices should be unsigned
+using Voxel = VoxelBase<ssize_t>;
+using Offset = VoxelBase<int>;
+
+} // namespace MR::Denoise::Kernel

From 915b1857a27845aa3cec3a88b52e91c29c358207 Mon Sep 17 00:00:00 2001
From: Robert Smith <robert.smith@florey.edu.au>
Date: Sun, 17 Nov 2024 22:09:34 +1100
Subject: [PATCH 18/34] New command dwi2noise

Closes #3035.
---
 cmd/dwi2noise.cpp                      | 194 +++++++++++++++++++++
 cmd/dwidenoise.cpp                     |  13 +-
 docs/reference/commands/dwi2noise.rst  | 123 +++++++++++++
 docs/reference/commands/dwidenoise.rst |  68 ++++++--
 docs/reference/commands_list.rst       |   2 +
 src/denoise/estimate.cpp               | 120 +++++++++++++
 src/denoise/{functor.h => estimate.h}  |  47 ++---
 src/denoise/functor.cpp                | 230 -------------------------
 src/denoise/recon.cpp                  | 163 ++++++++++++++++++
 src/denoise/recon.h                    |  67 +++++++
 10 files changed, 749 insertions(+), 278 deletions(-)
 create mode 100644 cmd/dwi2noise.cpp
 create mode 100644 docs/reference/commands/dwi2noise.rst
 create mode 100644 src/denoise/estimate.cpp
 rename src/denoise/{functor.h => estimate.h} (57%)
 delete mode 100644 src/denoise/functor.cpp
 create mode 100644 src/denoise/recon.cpp
 create mode 100644 src/denoise/recon.h

diff --git a/cmd/dwi2noise.cpp b/cmd/dwi2noise.cpp
new file mode 100644
index 0000000000..045dd5dd2a
--- /dev/null
+++ b/cmd/dwi2noise.cpp
@@ -0,0 +1,194 @@
+/* Copyright (c) 2008-2024 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#include <memory>
+
+#include "algo/threaded_loop.h"
+#include "command.h"
+#include "denoise/estimate.h"
+#include "denoise/estimator/estimator.h"
+#include "denoise/exports.h"
+#include "denoise/kernel/kernel.h"
+#include "exception.h"
+
+using namespace MR;
+using namespace App;
+using namespace MR::Denoise;
+
+// clang-format off
+void usage() {
+
+  SYNOPSIS = "Noise level estimation using Marchenko-Pastur PCA";
+
+  DESCRIPTION
+  + "DWI data noise map estimation"
+    " by interrogating data redundancy in the PCA domain"
+    " using the prior knowledge that the eigenspectrum of random covariance matrices"
+    " is described by the universal Marchenko-Pastur (MP) distribution."
+    " Fitting the MP distribution to the spectrum of patch-wise signal matrices"
+    " hence provides an estimator of the noise level 'sigma'."
+
+  + "Unlike the MRtrix3 command dwidenoise,"
+    " this command does not generate a denoised version of the input image series;"
+    " its primary output is instead a map of the estimated noise level."
+    " While this can also be obtained from the dwidenoise command using option -noise_out,"
+    " using instead the dwi2noise command gives the ability to obtain a noise map"
+    " to which filtering can be applied,"
+    " which can then be utilised for the actual image series denoising,"
+    " without generating an unwanted intermiedate denoised image series."
+
+  + "Important note:"
+    " noise level estimation should only be performed as the first step of an image processing pipeline."
+    " The routine is invalid if interpolation or smoothing has been applied to the data prior to denoising."
+
+  + "Note that on complex input data,"
+    " the output will be the total noise level across real and imaginary channels,"
+    " so a scale factor sqrt(2) applies."
+
+  + Kernel::shape_description
+
+  + Kernel::size_description;
+
+  AUTHOR = "Daan Christiaens (daan.christiaens@kcl.ac.uk)"
+           " and Jelle Veraart (jelle.veraart@nyumc.org)"
+           " and J-Donald Tournier (jdtournier@gmail.com)"
+           " and Robert E. Smith (robert.smith@florey.edu.au)";
+
+  REFERENCES
+  + "Veraart, J.; Fieremans, E. & Novikov, D.S. " // Internal
+    "Diffusion MRI noise mapping using random matrix theory. "
+    "Magn. Res. Med., 2016, 76(5), 1582-1593, doi: 10.1002/mrm.26059"
+
+  + "Cordero-Grande, L.; Christiaens, D.; Hutter, J.; Price, A.N.; Hajnal, J.V. " // Internal
+    "Complex diffusion-weighted image estimation via matrix recovery under general noise models. "
+    "NeuroImage, 2019, 200, 391-404, doi: 10.1016/j.neuroimage.2019.06.039"
+
+  + "* If using -estimator mrm2022: "
+    "Olesen, J.L.; Ianus, A.; Ostergaard, L.; Shemesh, N.; Jespersen, S.N. "
+    "Tensor denoising of multidimensional MRI data. "
+    "Magnetic Resonance in Medicine, 2022, 89(3), 1160-1172";
+
+  ARGUMENTS
+  + Argument("dwi", "the input diffusion-weighted image").type_image_in()
+  + Argument("noise", "the output estimated noise level map").type_image_out();
+
+  OPTIONS
+  + OptionGroup("Options for modifying PCA computations")
+  + datatype_option
+  + Estimator::option
+  + Kernel::options
+
+  // TODO Implement mask option
+  // Note that behaviour of -mask for dwi2noise may be different to that of dwidenoise
+
+  + OptionGroup("Options for exporting additional data regarding PCA behaviour")
+  + Option("rank",
+           "The signal rank estimated for the denoising patch centred at each input image voxel")
+    + Argument("image").type_image_out()
+  + OptionGroup("Options for debugging the operation of sliding window kernels")
+  + Option("max_dist",
+           "The maximum distance between a voxel and another voxel that was included in the local denoising patch")
+    + Argument("image").type_image_out()
+  + Option("voxelcount",
+           "The number of voxels that contributed to the PCA for processing of each voxel")
+    + Argument("image").type_image_out();
+
+  COPYRIGHT =
+      "Copyright (c) 2016 New York University, University of Antwerp, and the MRtrix3 contributors \n \n"
+      "Permission is hereby granted, free of charge, to any non-commercial entity ('Recipient') obtaining a copy of "
+      "this software and "
+      "associated documentation files (the 'Software'), to the Software solely for non-commercial research, including "
+      "the rights to "
+      "use, copy and modify the Software, subject to the following conditions: \n \n"
+      "\t 1. The above copyright notice and this permission notice shall be included by Recipient in all copies or "
+      "substantial portions of "
+      "the Software. \n \n"
+      "\t 2. THE SOFTWARE IS PROVIDED 'AS IS', WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT "
+      "LIMITED TO THE WARRANTIES"
+      "OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR "
+      "COPYRIGHT HOLDERS BE"
+      "LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING "
+      "FROM, OUT OF OR"
+      "IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. \n \n"
+      "\t 3. In no event shall NYU be liable for direct, indirect, special, incidental or consequential damages in "
+      "connection with the Software. "
+      "Recipient will defend, indemnify and hold NYU harmless from any claims or liability resulting from the use of "
+      "the Software by recipient. \n \n"
+      "\t 4. Neither anything contained herein nor the delivery of the Software to recipient shall be deemed to grant "
+      "the Recipient any right or "
+      "licenses under any patents or patent application owned by NYU. \n \n"
+      "\t 5. The Software may only be used for non-commercial research and may not be used for clinical care. \n \n"
+      "\t 6. Any publication by Recipient of research involving the Software shall cite the references listed below.";
+}
+// clang-format on
+
+template <typename T>
+void run(Header &data,
+         std::shared_ptr<Kernel::Base> kernel,
+         std::shared_ptr<Estimator::Base> estimator,
+         Exports &exports) {
+  auto input = data.get_image<T>().with_direct_io(3);
+  Image<bool> mask; // unused
+  Estimate<T> func(data, mask, kernel, estimator, exports);
+  ThreadedLoop("running MP-PCA noise level estimation", data, 0, 3).run(func, input);
+}
+
+void run() {
+  auto dwi = Header::open(argument[0]);
+
+  if (dwi.ndim() != 4 || dwi.size(3) <= 1)
+    throw Exception("input image must be 4-dimensional");
+
+  auto kernel = Kernel::make_kernel(dwi);
+  assert(kernel);
+
+  auto estimator = Estimator::make_estimator();
+  assert(estimator);
+
+  Exports exports(dwi);
+  exports.set_noise_out(argument[1]);
+  auto opt = get_options("rank");
+  if (!opt.empty())
+    exports.set_rank_input(opt[0][0]);
+  opt = get_options("max_dist");
+  if (!opt.empty())
+    exports.set_max_dist(opt[0][0]);
+  opt = get_options("voxelcount");
+  if (!opt.empty())
+    exports.set_voxelcount(opt[0][0]);
+
+  int prec = get_option_value("datatype", 0); // default: single precision
+  if (dwi.datatype().is_complex())
+    prec += 2; // support complex input data
+  switch (prec) {
+  case 0:
+    INFO("select real float32 for processing");
+    run<float>(dwi, kernel, estimator, exports);
+    break;
+  case 1:
+    INFO("select real float64 for processing");
+    run<double>(dwi, kernel, estimator, exports);
+    break;
+  case 2:
+    INFO("select complex float32 for processing");
+    run<cfloat>(dwi, kernel, estimator, exports);
+    break;
+  case 3:
+    INFO("select complex float64 for processing");
+    run<cdouble>(dwi, kernel, estimator, exports);
+    break;
+  }
+}
diff --git a/cmd/dwidenoise.cpp b/cmd/dwidenoise.cpp
index d90b0e1c7c..ad55fc03a9 100644
--- a/cmd/dwidenoise.cpp
+++ b/cmd/dwidenoise.cpp
@@ -14,7 +14,6 @@
  * For more details, see http://www.mrtrix.org/.
  */
 
-#include <mutex>
 #include <string>
 #include <vector>
 
@@ -32,12 +31,12 @@
 #include "denoise/estimator/mrm2022.h"
 #include "denoise/estimator/result.h"
 #include "denoise/exports.h"
-#include "denoise/functor.h"
 #include "denoise/kernel/cuboid.h"
 #include "denoise/kernel/data.h"
 #include "denoise/kernel/kernel.h"
 #include "denoise/kernel/sphere_radius.h"
 #include "denoise/kernel/sphere_ratio.h"
+#include "denoise/recon.h"
 
 using namespace MR;
 using namespace App;
@@ -72,7 +71,7 @@ void usage() {
 
   + "By default, optimal value shrinkage based on minimisation of the Frobenius norm "
     "will be used to attenuate eigenvectors based on the estimated noise level. "
-    "Hard truncation of sub-threshold components"
+    "Hard truncation of sub-threshold components and inclusion of supra-threshold components"
     "---which was the behaviour of the dwidenoise command in version 3.0.x---"
     "can be activated using -filter truncate."
 
@@ -85,9 +84,9 @@ void usage() {
     "There are multiple algebraic forms that modulate the weight with which each decomposition "
     "contributes with greater or lesser strength toward the output image intensities. "
     "The various options are: "
-    "'Gaussian': A Gaussian distribution with FWHM equal to twice the voxel size, "
+    "'gaussian': A Gaussian distribution with FWHM equal to twice the voxel size, "
       "such that decompisitions centred more closely to the output voxel have greater influence; "
-    "'invL0': The inverse of the L0 norm (ie. rank) of each decomposition, "
+    "'invl0': The inverse of the L0 norm (ie. rank) of each decomposition, "
       "as used in Manjon et al. 2013; "
     "'rank': The rank of each decomposition, "
       "such that high-rank decompositions contribute more strongly to the output intensities "
@@ -130,10 +129,10 @@ void usage() {
   + OptionGroup("Options for modifying PCA computations")
   + datatype_option
   + Estimator::option
-
   + Kernel::options
 
   + OptionGroup("Options that affect reconstruction of the output image series")
+  // TODO Separate masks for voxels to contribute to patches vs. voxels for which to perform denoising
   + Option("mask",
            "Only denoise voxels within the specified binary brain mask image.")
     + Argument("image").type_image_in()
@@ -249,7 +248,7 @@ void run(Header &data,
   header.datatype() = DataType::from<T>();
   auto output = Image<T>::create(output_name, header);
   // run
-  Functor<T> func(data, mask, kernel, estimator, filter, aggregator, exports);
+  Recon<T> func(data, mask, kernel, estimator, filter, aggregator, exports);
   ThreadedLoop("running MP-PCA denoising", data, 0, 3).run(func, input, output);
   // Rescale output if performing aggregation
   if (aggregator == aggregator_type::EXCLUSIVE)
diff --git a/docs/reference/commands/dwi2noise.rst b/docs/reference/commands/dwi2noise.rst
new file mode 100644
index 0000000000..cf4a1f1383
--- /dev/null
+++ b/docs/reference/commands/dwi2noise.rst
@@ -0,0 +1,123 @@
+.. _dwi2noise:
+
+dwi2noise
+===================
+
+Synopsis
+--------
+
+Noise level estimation using Marchenko-Pastur PCA
+
+Usage
+--------
+
+::
+
+    dwi2noise [ options ]  dwi noise
+
+-  *dwi*: the input diffusion-weighted image
+-  *noise*: the output estimated noise level map
+
+Description
+-----------
+
+DWI data noise map estimation by exploiting data redundancy in the PCA domain using the prior knowledge that the eigenspectrum of random covariance matrices is described by the universal Marchenko-Pastur (MP) distribution. Fitting the MP distribution to the spectrum of patch-wise signal matrices hence provides an estimator of the noise level 'sigma'.
+
+Unlike the MRtrix3 command dwidenoise, this command does not generate a denoised version of the input image series; its primary output is instead a map of the estimated noise level. While this can also be obtained from the dwidenoise command using option -noise_out, using instead the dwi2noise command gives the ability to obtain a noise map to which filtering can be applied, which can then be utilised for the actual image series denoising, without generating an unwanted intermiedate denoised image series.
+
+Important note: noise level estimation should only be performed as the first step of an image processing pipeline. The routine is invalid if interpolation or smoothing has been applied to the data prior to denoising.
+
+Note that on complex input data, the output will be the total noise level across real and imaginary channels, so a scale factor sqrt(2) applies.
+
+The sliding spatial window behaves differently at the edges of the image FoV depending on the shape / size selected for that window. The default behaviour is to use a spherical kernel centred at the voxel of interest, whose size is some multiple of the number of input volumes; where some such voxels lie outside of the image FoV, the radius of the kernel will be increased until the requisite number of voxels are used. For a spherical kernel of a fixed radius, no such expansion will occur, and so for voxels near the image edge a reduced number of voxels will be present in the kernel. For a cuboid kernel, the centre of the kernel will be offset from the voxel being processed such that the entire volume of the kernel resides within the image FoV.
+
+The size of the default spherical kernel is set to select a number of voxels that is 1.0 / 0.85 ~ 1.18 times the number of volumes in the input series. If a cuboid kernel is requested, but the -extent option is not specified, the command will select the smallest isotropic patch size that exceeds the number of DW images in the input data; e.g., 5x5x5 for data with <= 125 DWI volumes, 7x7x7 for data with <= 343 DWI volumes, etc.
+
+Options
+-------
+
+Options for modifying PCA computations
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+-  **-datatype float32/float64** Datatype for the eigenvalue decomposition (single or double precision). For complex input data, this will select complex float32 or complex float64 datatypes.
+
+-  **-estimator algorithm** Select the noise level estimator (default = Exp2), either:  |br|
+   * Exp1: the original estimator used in Veraart et al. (2016);  |br|
+   * Exp2: the improved estimator introduced in Cordero-Grande et al. (2019);  |br|
+   * MRM2022: the alternative estimator introduced in Olesen et al. (2022).
+
+Options for controlling the sliding spatial window kernel
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+-  **-shape choice** Set the shape of the sliding spatial window. Options are: cuboid,sphere; default: sphere
+
+-  **-radius_mm value** Set an absolute spherical kernel radius in mm
+
+-  **-radius_ratio value** Set the spherical kernel size as a ratio of number of voxels to number of input volumes (default: 1.0/0.85 ~= 1.18)
+
+-  **-extent window** Set the patch size of the cuboid kernel; can be either a single odd integer or a comma-separated triplet of odd integers
+
+Options for exporting additional data regarding PCA behaviour
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+-  **-rank image** The signal rank estimated for the denoising patch centred at each input image voxel
+
+Options for debugging the operation of sliding window kernels
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+-  **-max_dist image** The maximum distance between a voxel and another voxel that was included in the local denoising patch
+
+-  **-voxelcount image** The number of voxels that contributed to the PCA for processing of each voxel
+
+Standard options
+^^^^^^^^^^^^^^^^
+
+-  **-info** display information messages.
+
+-  **-quiet** do not display information messages or progress status; alternatively, this can be achieved by setting the MRTRIX_QUIET environment variable to a non-empty string.
+
+-  **-debug** display debugging messages.
+
+-  **-force** force overwrite of output files (caution: using the same file as input and output might cause unexpected behaviour).
+
+-  **-nthreads number** use this number of threads in multi-threaded applications (set to 0 to disable multi-threading).
+
+-  **-config key value** *(multiple uses permitted)* temporarily set the value of an MRtrix config file entry.
+
+-  **-help** display this information page and exit.
+
+-  **-version** display version information and exit.
+
+References
+^^^^^^^^^^
+
+Veraart, J.; Fieremans, E. & Novikov, D.S. Diffusion MRI noise mapping using random matrix theory. Magn. Res. Med., 2016, 76(5), 1582-1593, doi: 10.1002/mrm.26059
+
+Cordero-Grande, L.; Christiaens, D.; Hutter, J.; Price, A.N.; Hajnal, J.V. Complex diffusion-weighted image estimation via matrix recovery under general noise models. NeuroImage, 2019, 200, 391-404, doi: 10.1016/j.neuroimage.2019.06.039
+
+* If using -estimator mrm2022: Olesen, J.L.; Ianus, A.; Ostergaard, L.; Shemesh, N.; Jespersen, S.N. Tensor denoising of multidimensional MRI data. Magnetic Resonance in Medicine, 2022, 89(3), 1160-1172
+
+Tournier, J.-D.; Smith, R. E.; Raffelt, D.; Tabbara, R.; Dhollander, T.; Pietsch, M.; Christiaens, D.; Jeurissen, B.; Yeh, C.-H. & Connelly, A. MRtrix3: A fast, flexible and open software framework for medical image processing and visualisation. NeuroImage, 2019, 202, 116137
+
+--------------
+
+
+
+**Author:** Daan Christiaens (daan.christiaens@kcl.ac.uk) and Jelle Veraart (jelle.veraart@nyumc.org) and J-Donald Tournier (jdtournier@gmail.com) and Robert E. Smith (robert.smith@florey.edu.au)
+
+**Copyright:** Copyright (c) 2016 New York University, University of Antwerp, and the MRtrix3 contributors 
+ 
+Permission is hereby granted, free of charge, to any non-commercial entity ('Recipient') obtaining a copy of this software and associated documentation files (the 'Software'), to the Software solely for non-commercial research, including the rights to use, copy and modify the Software, subject to the following conditions: 
+ 
+	 1. The above copyright notice and this permission notice shall be included by Recipient in all copies or substantial portions of the Software. 
+ 
+	 2. THE SOFTWARE IS PROVIDED 'AS IS', WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIESOF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BELIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF ORIN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 
+ 
+	 3. In no event shall NYU be liable for direct, indirect, special, incidental or consequential damages in connection with the Software. Recipient will defend, indemnify and hold NYU harmless from any claims or liability resulting from the use of the Software by recipient. 
+ 
+	 4. Neither anything contained herein nor the delivery of the Software to recipient shall be deemed to grant the Recipient any right or licenses under any patents or patent application owned by NYU. 
+ 
+	 5. The Software may only be used for non-commercial research and may not be used for clinical care. 
+ 
+	 6. Any publication by Recipient of research involving the Software shall cite the references listed below.
+
diff --git a/docs/reference/commands/dwidenoise.rst b/docs/reference/commands/dwidenoise.rst
index a1a4660b08..c769924859 100644
--- a/docs/reference/commands/dwidenoise.rst
+++ b/docs/reference/commands/dwidenoise.rst
@@ -21,28 +21,72 @@ Usage
 Description
 -----------
 
-DWI data denoising and noise map estimation by exploiting data redundancy in the PCA domain using the prior knowledge that the eigenspectrum of random covariance matrices is described by the universal Marchenko-Pastur (MP) distribution. Fitting the MP distribution to the spectrum of patch-wise signal matrices hence provides an estimator of the noise level 'sigma', as was first shown in Veraart et al. (2016) and later improved in Cordero-Grande et al. (2019). This noise level estimate then determines the optimal cut-off for PCA denoising.
+DWI data denoising and noise map estimation by exploiting data redundancy in the PCA domain using the prior knowledge that the eigenspectrum of random covariance matrices is described by the universal Marchenko-Pastur (MP) distribution. Fitting the MP distribution to the spectrum of patch-wise signal matrices hence provides an estimator of the noise level 'sigma'; this noise level estimate then determines the optimal cut-off for PCA denoising.
 
 Important note: image denoising must be performed as the first step of the image processing pipeline. The routine will fail if interpolation or smoothing has been applied to the data prior to denoising.
 
 Note that this function does not correct for non-Gaussian noise biases present in magnitude-reconstructed MRI images. If available, including the MRI phase data can reduce such non-Gaussian biases, and the command now supports complex input data.
 
+The sliding spatial window behaves differently at the edges of the image FoV depending on the shape / size selected for that window. The default behaviour is to use a spherical kernel centred at the voxel of interest, whose size is some multiple of the number of input volumes; where some such voxels lie outside of the image FoV, the radius of the kernel will be increased until the requisite number of voxels are used. For a spherical kernel of a fixed radius, no such expansion will occur, and so for voxels near the image edge a reduced number of voxels will be present in the kernel. For a cuboid kernel, the centre of the kernel will be offset from the voxel being processed such that the entire volume of the kernel resides within the image FoV.
+
+The size of the default spherical kernel is set to select a number of voxels that is 1.0 / 0.85 ~ 1.18 times the number of volumes in the input series. If a cuboid kernel is requested, but the -extent option is not specified, the command will select the smallest isotropic patch size that exceeds the number of DW images in the input data; e.g., 5x5x5 for data with <= 125 DWI volumes, 7x7x7 for data with <= 343 DWI volumes, etc.
+
+By default, optimal value shrinkage based on minimisation of the Frobenius norm will be used to attenuate eigenvectors based on the estimated noise level. Hard truncation of sub-threshold components and inclusion of supra-threshold components---which was the behaviour of the dwidenoise command in version 3.0.x---can be activated using -filter truncate.
+
+-aggregation exclusive corresponds to the behaviour of the dwidenoise command in version 3.0.x, where the output intensities for a given image voxel are determined exclusively from the PCA decomposition where the sliding spatial window is centred at that voxel. In all other use cases, so-called "overcomplete local PCA" is performed, where the intensities for an output image voxel are some combination of all PCA decompositions for which that voxel is included in the local spatial kernel. There are multiple algebraic forms that modulate the weight with which each decomposition contributes with greater or lesser strength toward the output image intensities. The various options are: 'gaussian': A Gaussian distribution with FWHM equal to twice the voxel size, such that decompisitions centred more closely to the output voxel have greater influence; 'invl0': The inverse of the L0 norm (ie. rank) of each decomposition, as used in Manjon et al. 2013; 'rank': The rank of each decomposition, such that high-rank decompositions contribute more strongly to the output intensities regardless of distance between the output voxel and the centre of the decomposition kernel; 'uniform': All decompositions that include the output voxel in the sliding spatial window contribute equally.
+
 Options
 -------
 
--  **-mask image** Only process voxels within the specified binary brain mask image.
+Options for modifying PCA computations
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
--  **-extent window** Set the patch size of the denoising filter. By default, the command will select the smallest isotropic patch size that exceeds the number of DW images in the input data, e.g., 5x5x5 for data with <= 125 DWI volumes, 7x7x7 for data with <= 343 DWI volumes, etc.
+-  **-datatype float32/float64** Datatype for the eigenvalue decomposition (single or double precision). For complex input data, this will select complex float32 or complex float64 datatypes.
 
--  **-noise level** The output noise map, i.e., the estimated noise level 'sigma' in the data.Note that on complex input data, this will be the total noise level across real and imaginary channels, so a scale factor sqrt(2) applies.
+-  **-estimator algorithm** Select the noise level estimator (default = Exp2), either:  |br|
+   * Exp1: the original estimator used in Veraart et al. (2016);  |br|
+   * Exp2: the improved estimator introduced in Cordero-Grande et al. (2019);  |br|
+   * MRM2022: the alternative estimator introduced in Olesen et al. (2022).
 
--  **-rank cutoff** The selected signal rank of the output denoised image.
+Options for controlling the sliding spatial window kernel
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
--  **-datatype float32/float64** Datatype for the eigenvalue decomposition (single or double precision). For complex input data, this will select complex float32 or complex float64 datatypes.
+-  **-shape choice** Set the shape of the sliding spatial window. Options are: cuboid,sphere; default: sphere
+
+-  **-radius_mm value** Set an absolute spherical kernel radius in mm
+
+-  **-radius_ratio value** Set the spherical kernel size as a ratio of number of voxels to number of input volumes (default: 1.0/0.85 ~= 1.18)
+
+-  **-extent window** Set the patch size of the cuboid kernel; can be either a single odd integer or a comma-separated triplet of odd integers
+
+Options that affect reconstruction of the output image series
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+-  **-mask image** Only denoise voxels within the specified binary brain mask image.
+
+-  **-filter choice** Modulate how component contributions are filtered based on the cumulative eigenvalues relative to the noise level; options are: truncate,frobenius; default: frobenius (Optimal Shrinkage based on minimisation of the Frobenius norm)
 
--  **-estimator Exp1/Exp2** Select the noise level estimator (default = Exp2), either:  |br|
-   * Exp1: the original estimator used in Veraart et al. (2016), or  |br|
-   * Exp2: the improved estimator introduced in Cordero-Grande et al. (2019).
+-  **-aggregator choice** Select how the outcomes of multiple PCA outcomes centred at different voxels contribute to the reconstructed DWI signal in each voxel; options are: exclusive,gaussian,invl0,rank,uniform; default: Gaussian
+
+Options for exporting additional data regarding PCA behaviour
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+-  **-noise_out image** The output noise map, i.e., the estimated noise level 'sigma' in the data. Note that on complex input data, this will be the total noise level across real and imaginary channels, so a scale factor sqrt(2) applies.
+
+-  **-rank_input image** The signal rank estimated for the denoising patch centred at each input image voxel
+
+-  **-rank_output image** An estimated rank for the output image data, accounting for multi-patch aggregation
+
+Options for debugging the operation of sliding window kernels
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+-  **-max_dist image** The maximum distance between a voxel and another voxel that was included in the local denoising patch
+
+-  **-voxelcount image** The number of voxels that contributed to the PCA for processing of each voxel
+
+-  **-sum_aggregation image** The sum of aggregation weights of those patches contributing to each output voxel
+
+-  **-sum_optshrink image** the sum of eigenvector weights computed for the denoising patch centred at each voxel as a result of performing optimal shrinkage
 
 Standard options
 ^^^^^^^^^^^^^^^^
@@ -72,13 +116,17 @@ Veraart, J.; Fieremans, E. & Novikov, D.S. Diffusion MRI noise mapping using ran
 
 Cordero-Grande, L.; Christiaens, D.; Hutter, J.; Price, A.N.; Hajnal, J.V. Complex diffusion-weighted image estimation via matrix recovery under general noise models. NeuroImage, 2019, 200, 391-404, doi: 10.1016/j.neuroimage.2019.06.039
 
+* If using -estimator mrm2022: Olesen, J.L.; Ianus, A.; Ostergaard, L.; Shemesh, N.; Jespersen, S.N. Tensor denoising of multidimensional MRI data. Magnetic Resonance in Medicine, 2022, 89(3), 1160-1172
+
+* If using anything other than -aggregation exclusive: Manjon, J.V.; Coupe, P.; Concha, L.; Buades, A.; D. Collins, D.L.; Robles, M. Diffusion Weighted Image Denoising Using Overcomplete Local PCA. PLoS ONE, 2013, 8(9), e73021
+
 Tournier, J.-D.; Smith, R. E.; Raffelt, D.; Tabbara, R.; Dhollander, T.; Pietsch, M.; Christiaens, D.; Jeurissen, B.; Yeh, C.-H. & Connelly, A. MRtrix3: A fast, flexible and open software framework for medical image processing and visualisation. NeuroImage, 2019, 202, 116137
 
 --------------
 
 
 
-**Author:** Daan Christiaens (daan.christiaens@kcl.ac.uk) and Jelle Veraart (jelle.veraart@nyumc.org) and J-Donald Tournier (jdtournier@gmail.com)
+**Author:** Daan Christiaens (daan.christiaens@kcl.ac.uk) and Jelle Veraart (jelle.veraart@nyumc.org) and J-Donald Tournier (jdtournier@gmail.com) and Robert E. Smith (robert.smith@florey.edu.au)
 
 **Copyright:** Copyright (c) 2016 New York University, University of Antwerp, and the MRtrix3 contributors 
  
diff --git a/docs/reference/commands_list.rst b/docs/reference/commands_list.rst
index 78cfaa506e..b32465b80e 100644
--- a/docs/reference/commands_list.rst
+++ b/docs/reference/commands_list.rst
@@ -34,6 +34,7 @@ List of MRtrix3 commands
     commands/dwi2adc
     commands/dwi2fod
     commands/dwi2mask
+    commands/dwi2noise
     commands/dwi2response
     commands/dwi2tensor
     commands/dwibiascorrect
@@ -167,6 +168,7 @@ List of MRtrix3 commands
     |cpp.png|, :ref:`dwi2adc`, "Calculate ADC and/or IVIM parameters."
     |cpp.png|, :ref:`dwi2fod`, "Estimate fibre orientation distributions from diffusion data using spherical deconvolution"
     |python.png|, :ref:`dwi2mask`, "Generate a binary mask from DWI data"
+    |cpp.png|, :ref:`dwi2noise`, "Noise level estimation using Marchenko-Pastur PCA"
     |python.png|, :ref:`dwi2response`, "Estimate response function(s) for spherical deconvolution"
     |cpp.png|, :ref:`dwi2tensor`, "Diffusion (kurtosis) tensor estimation"
     |python.png|, :ref:`dwibiascorrect`, "Perform B1 field inhomogeneity correction for a DWI volume series"
diff --git a/src/denoise/estimate.cpp b/src/denoise/estimate.cpp
new file mode 100644
index 0000000000..c2c19a269a
--- /dev/null
+++ b/src/denoise/estimate.cpp
@@ -0,0 +1,120 @@
+/* Copyright (c) 2008-2024 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#include "denoise/estimate.h"
+
+#include <limits>
+
+#include "math/math.h"
+
+namespace MR::Denoise {
+
+template <typename F>
+Estimate<F>::Estimate(const Header &header,
+                      Image<bool> &mask,
+                      std::shared_ptr<Kernel::Base> kernel,
+                      std::shared_ptr<Estimator::Base> estimator,
+                      Exports &exports)
+    : m(header.size(3)),
+      mask(mask),
+      kernel(kernel),
+      estimator(estimator),
+      X(m, kernel->estimated_size()),
+      XtX(std::min(m, kernel->estimated_size()), std::min(m, kernel->estimated_size())),
+      eig(std::min(m, kernel->estimated_size())),
+      s(std::min(m, kernel->estimated_size())),
+      exports(exports) {}
+
+template <typename F> void Estimate<F>::operator()(Image<F> &dwi) {
+
+  // Process voxels in mask only
+  if (mask.valid()) {
+    assign_pos_of(dwi, 0, 3).to(mask);
+    if (!mask.value())
+      return;
+  }
+
+  // Load list of voxels from which to load data
+  neighbourhood = (*kernel)({dwi.index(0), dwi.index(1), dwi.index(2)});
+  const ssize_t n = neighbourhood.voxels.size();
+  const ssize_t r = std::min(m, n);
+  const ssize_t q = std::max(m, n);
+
+  // Expand local storage if necessary
+  if (n > X.cols()) {
+    DEBUG("Expanding data matrix storage from " + str(m) + "x" + str(X.cols()) + " to " + str(m) + "x" + str(n));
+    X.resize(m, n);
+  }
+  if (r > XtX.cols()) {
+    DEBUG("Expanding decomposition matrix storage from " + str(X.rows()) + " to " + str(r));
+    XtX.resize(r, r);
+    s.resize(r);
+  }
+
+  // Fill matrices with NaN when in debug mode;
+  //   make sure results from one voxel are not creeping into another
+  //   due to use of block oberations to prevent memory re-allocation
+  //   in the presence of variation in kernel sizes
+#ifndef NDEBUG
+  X.fill(std::numeric_limits<F>::signaling_NaN());
+  XtX.fill(std::numeric_limits<F>::signaling_NaN());
+  s.fill(std::numeric_limits<default_type>::signaling_NaN());
+#endif
+
+  load_data(dwi, neighbourhood.voxels);
+
+  // Compute Eigendecomposition:
+  if (m <= n)
+    XtX.topLeftCorner(r, r).template triangularView<Eigen::Lower>() = X.leftCols(n) * X.leftCols(n).adjoint();
+  else
+    XtX.topLeftCorner(r, r).template triangularView<Eigen::Lower>() = X.leftCols(n).adjoint() * X.leftCols(n);
+  eig.compute(XtX.topLeftCorner(r, r));
+  // eigenvalues sorted in increasing order:
+  s.head(r) = eig.eigenvalues().template cast<double>();
+
+  // Marchenko-Pastur optimal threshold determination
+  threshold = (*estimator)(s, m, n);
+  const ssize_t in_rank = r - threshold.cutoff_p;
+
+  // Store additional output maps if requested
+  if (exports.noise_out.valid()) {
+    assign_pos_of(dwi, 0, 3).to(exports.noise_out);
+    exports.noise_out.value() = float(std::sqrt(threshold.sigma2));
+  }
+  if (exports.rank_input.valid()) {
+    assign_pos_of(dwi, 0, 3).to(exports.rank_input);
+    exports.rank_input.value() = in_rank;
+  }
+  if (exports.max_dist.valid()) {
+    assign_pos_of(dwi, 0, 3).to(exports.max_dist);
+    exports.max_dist.value() = neighbourhood.max_distance;
+  }
+  if (exports.voxelcount.valid()) {
+    assign_pos_of(dwi, 0, 3).to(exports.voxelcount);
+    exports.voxelcount.value() = n;
+  }
+}
+
+template <typename F> void Estimate<F>::load_data(Image<F> &image, const std::vector<Kernel::Voxel> &voxels) {
+  const Kernel::Voxel::index_type pos({image.index(0), image.index(1), image.index(2)});
+  for (ssize_t i = 0; i != voxels.size(); ++i) {
+    assign_pos_of(voxels[i].index, 0, 3).to(image);
+    X.col(i) = image.row(3);
+  }
+  assign_pos_of(pos, 0, 3).to(image);
+}
+
+} // namespace MR::Denoise
diff --git a/src/denoise/functor.h b/src/denoise/estimate.h
similarity index 57%
rename from src/denoise/functor.h
rename to src/denoise/estimate.h
index 11cf94c42a..babbabfb60 100644
--- a/src/denoise/functor.h
+++ b/src/denoise/estimate.h
@@ -16,9 +16,7 @@
 
 #pragma once
 
-#include <limits>
 #include <memory>
-#include <mutex>
 
 #include <Eigen/Dense>
 
@@ -34,58 +32,45 @@
 
 namespace MR::Denoise {
 
-template <typename F> class Functor {
+template <typename F> class Estimate {
 
 public:
   using MatrixType = Eigen::Matrix<F, Eigen::Dynamic, Eigen::Dynamic>;
 
-  Functor(const Header &header,
-          Image<bool> &mask,
-          std::shared_ptr<Kernel::Base> kernel,
-          std::shared_ptr<Estimator::Base> estimator,
-          filter_type filter,
-          aggregator_type aggregator,
-          Exports &exports);
+  Estimate(const Header &header,
+           Image<bool> &mask,
+           std::shared_ptr<Kernel::Base> kernel,
+           std::shared_ptr<Estimator::Base> estimator,
+           Exports &exports);
 
-  void operator()(Image<F> &dwi, Image<F> &out);
+  void operator()(Image<F> &dwi);
 
-private:
-  // Denoising configuration
+protected:
   const ssize_t m;
+
+  // Denoising configuration
   Image<bool> mask;
   std::shared_ptr<Kernel::Base> kernel;
   std::shared_ptr<Estimator::Base> estimator;
-  filter_type filter;
-  aggregator_type aggregator;
-  double gaussian_multiplier;
 
   // Reusable memory
+  Kernel::Data neighbourhood;
   MatrixType X;
   MatrixType XtX;
   Eigen::SelfAdjointEigenSolver<MatrixType> eig;
   eigenvalues_type s;
+  Estimator::Result threshold;
   vector_type clam;
-  vector_type w;
 
   // Export images
   Exports exports;
 
-  // Data that can only be written in a thread-safe manner
-  // Note that this applies not just to this scratch buffer, but also the output image
-  //   (while it would be thread-safe to create a full copy of the output image for each thread
-  //   and combine them only at destruction time,
-  //   this runs the risk of becoming prohibitively large)
-  // Not placing this within a MutexProtexted<> as the image type is still templated
-  static std::mutex mutex_aggregator;
-
   void load_data(Image<F> &image, const std::vector<Kernel::Voxel> &voxels);
 };
 
-template <typename F> std::mutex Functor<F>::mutex_aggregator;
-
-template class Functor<float>;
-template class Functor<cfloat>;
-template class Functor<double>;
-template class Functor<cdouble>;
+template class Estimate<float>;
+template class Estimate<cfloat>;
+template class Estimate<double>;
+template class Estimate<cdouble>;
 
 } // namespace MR::Denoise
diff --git a/src/denoise/functor.cpp b/src/denoise/functor.cpp
deleted file mode 100644
index 4fab62d51f..0000000000
--- a/src/denoise/functor.cpp
+++ /dev/null
@@ -1,230 +0,0 @@
-/* Copyright (c) 2008-2024 the MRtrix3 contributors.
- *
- * This Source Code Form is subject to the terms of the Mozilla Public
- * License, v. 2.0. If a copy of the MPL was not distributed with this
- * file, You can obtain one at http://mozilla.org/MPL/2.0/.
- *
- * Covered Software is provided under this License on an "as is"
- * basis, without warranty of any kind, either expressed, implied, or
- * statutory, including, without limitation, warranties that the
- * Covered Software is free of defects, merchantable, fit for a
- * particular purpose or non-infringing.
- * See the Mozilla Public License v. 2.0 for more details.
- *
- * For more details, see http://www.mrtrix.org/.
- */
-
-#include "denoise/functor.h"
-
-#include "math/math.h"
-
-namespace MR::Denoise {
-
-template <typename F>
-Functor<F>::Functor(const Header &header,
-                    Image<bool> &mask,
-                    std::shared_ptr<Kernel::Base> kernel,
-                    std::shared_ptr<Estimator::Base> estimator,
-                    filter_type filter,
-                    aggregator_type aggregator,
-                    Exports &exports)
-    : m(header.size(3)),
-      mask(mask),
-      kernel(kernel),
-      estimator(estimator),
-      filter(filter),
-      aggregator(aggregator),
-      // FWHM = 2 x cube root of voxel spacings
-      gaussian_multiplier(-std::log(2.0) /
-                          Math::pow2(std::cbrt(header.spacing(0) * header.spacing(1) * header.spacing(2)))),
-      X(m, kernel->estimated_size()),
-      XtX(std::min(m, kernel->estimated_size()), std::min(m, kernel->estimated_size())),
-      eig(std::min(m, kernel->estimated_size())),
-      s(std::min(m, kernel->estimated_size())),
-      clam(std::min(m, kernel->estimated_size())),
-      w(std::min(m, kernel->estimated_size())),
-      exports(exports) {}
-
-template <typename F> void Functor<F>::operator()(Image<F> &dwi, Image<F> &out) {
-  // Process voxels in mask only
-  if (mask.valid()) {
-    assign_pos_of(dwi, 0, 3).to(mask);
-    if (!mask.value())
-      return;
-  }
-
-  // Load list of voxels from which to load data
-  const Kernel::Data neighbourhood = (*kernel)({dwi.index(0), dwi.index(1), dwi.index(2)});
-  const ssize_t n = neighbourhood.voxels.size();
-  const ssize_t r = std::min(m, n);
-  const ssize_t q = std::max(m, n);
-
-  // Expand local storage if necessary
-  if (n > X.cols()) {
-    DEBUG("Expanding data matrix storage from " + str(m) + "x" + str(X.cols()) + " to " + str(m) + "x" + str(n));
-    X.resize(m, n);
-  }
-  if (r > XtX.cols()) {
-    DEBUG("Expanding decomposition matrix storage from " + str(X.rows()) + " to " + str(r));
-    XtX.resize(r, r);
-    s.resize(r);
-    clam.resize(r);
-    w.resize(r);
-  }
-
-  // Fill matrices with NaN when in debug mode;
-  //   make sure results from one voxel are not creeping into another
-  //   due to use of block oberations to prevent memory re-allocation
-  //   in the presence of variation in kernel sizes
-#ifndef NDEBUG
-  X.fill(std::numeric_limits<F>::signaling_NaN());
-  XtX.fill(std::numeric_limits<F>::signaling_NaN());
-  s.fill(std::numeric_limits<default_type>::signaling_NaN());
-  clam.fill(std::numeric_limits<default_type>::signaling_NaN());
-  w.fill(std::numeric_limits<default_type>::signaling_NaN());
-#endif
-
-  load_data(dwi, neighbourhood.voxels);
-
-  // Compute Eigendecomposition:
-  if (m <= n)
-    XtX.topLeftCorner(r, r).template triangularView<Eigen::Lower>() = X.leftCols(n) * X.leftCols(n).adjoint();
-  else
-    XtX.topLeftCorner(r, r).template triangularView<Eigen::Lower>() = X.leftCols(n).adjoint() * X.leftCols(n);
-  eig.compute(XtX.topLeftCorner(r, r));
-  // eigenvalues sorted in increasing order:
-  s.head(r) = eig.eigenvalues().template cast<double>();
-
-  // Marchenko-Pastur optimal threshold determination
-  const Estimator::Result threshold = (*estimator)(s, m, n);
-
-  // Generate weights vector
-  double sum_weights = 0.0;
-  ssize_t out_rank = 0;
-  switch (filter) {
-  case filter_type::TRUNCATE:
-    out_rank = r - threshold.cutoff_p;
-    w.head(threshold.cutoff_p).setZero();
-    w.segment(threshold.cutoff_p, r - threshold.cutoff_p).setOnes();
-    sum_weights = double(out_rank);
-    break;
-  case filter_type::FROBENIUS: {
-    const double beta = r / q;
-    const double transition = 1.0 + std::sqrt(beta);
-    double clam = 0.0;
-    for (ssize_t i = 0; i != r; ++i) {
-      const double lam = std::max(s[i], 0.0) / q;
-      clam += lam;
-      const double y = clam / (threshold.sigma2 * (i + 1));
-      double nu = 0.0;
-      if (y > transition) {
-        nu = std::sqrt(Math::pow2(Math::pow2(y) - beta - 1.0) - (4.0 * beta)) / y;
-        ++out_rank;
-      }
-      w[i] = nu / y;
-      sum_weights += w[i];
-    }
-  } break;
-  default:
-    assert(false);
-  }
-
-  // recombine data using only eigenvectors above threshold
-  // If only the data computed when this voxel was the centre of the patch
-  //   is to be used for synthesis of the output image,
-  //   then only that individual column needs to be reconstructed;
-  //   if however the result from this patch is to contribute to the synthesized image
-  //   for all voxels that were utilised within this patch,
-  //   then we need to instead compute the full projection
-  switch (aggregator) {
-  case aggregator_type::EXCLUSIVE:
-    if (m <= n)
-      X.col(neighbourhood.centre_index) =
-          eig.eigenvectors() * (w.head(r).cast<F>().matrix().asDiagonal() *
-                                (eig.eigenvectors().adjoint() * X.col(neighbourhood.centre_index)));
-    else
-      X.col(neighbourhood.centre_index) =
-          X.leftCols(n) * (eig.eigenvectors() * (w.head(r).cast<F>().matrix().asDiagonal() *
-                                                 eig.eigenvectors().adjoint().col(neighbourhood.centre_index)));
-    assign_pos_of(dwi).to(out);
-    out.row(3) = X.col(neighbourhood.centre_index);
-    if (exports.sum_aggregation.valid()) {
-      assign_pos_of(dwi, 0, 3).to(exports.sum_aggregation);
-      exports.sum_aggregation.value() = 1.0;
-    }
-    if (exports.rank_output.valid()) {
-      assign_pos_of(dwi, 0, 3).to(exports.rank_output);
-      exports.rank_output.value() = out_rank;
-    }
-    break;
-  default: {
-    if (m <= n)
-      X = eig.eigenvectors() * (w.head(r).cast<F>().matrix().asDiagonal() * (eig.eigenvectors().adjoint() * X));
-    else
-      X.leftCols(n) = X.leftCols(n) *
-                      (eig.eigenvectors() * (w.head(r).cast<F>().matrix().asDiagonal() * eig.eigenvectors().adjoint()));
-    std::lock_guard<std::mutex> lock(mutex_aggregator);
-    for (size_t voxel_index = 0; voxel_index != neighbourhood.voxels.size(); ++voxel_index) {
-      assign_pos_of(neighbourhood.voxels[voxel_index].index, 0, 3).to(out);
-      assign_pos_of(neighbourhood.voxels[voxel_index].index).to(exports.sum_aggregation);
-      double weight = std::numeric_limits<double>::signaling_NaN();
-      switch (aggregator) {
-      case aggregator_type::EXCLUSIVE:
-        assert(false);
-        break;
-      case aggregator_type::GAUSSIAN:
-        weight = std::exp(gaussian_multiplier * neighbourhood.voxels[voxel_index].sq_distance);
-        break;
-      case aggregator_type::INVL0:
-        weight = 1.0 / (1 + out_rank);
-        break;
-      case aggregator_type::RANK:
-        weight = out_rank;
-        break;
-      case aggregator_type::UNIFORM:
-        weight = 1.0;
-        break;
-      }
-      out.row(3) += weight * X.col(voxel_index);
-      exports.sum_aggregation.value() += weight;
-      if (exports.rank_output.valid()) {
-        assign_pos_of(neighbourhood.voxels[voxel_index].index, 0, 3).to(exports.rank_output);
-        exports.rank_output.value() += weight * out_rank;
-      }
-    }
-  } break;
-  }
-
-  // Store additional output maps if requested
-  if (exports.noise_out.valid()) {
-    assign_pos_of(dwi, 0, 3).to(exports.noise_out);
-    exports.noise_out.value() = float(std::sqrt(threshold.sigma2));
-  }
-  if (exports.rank_input.valid()) {
-    assign_pos_of(dwi, 0, 3).to(exports.rank_input);
-    exports.rank_input.value() = out_rank;
-  }
-  if (exports.sum_optshrink.valid()) {
-    assign_pos_of(dwi, 0, 3).to(exports.sum_optshrink);
-    exports.sum_optshrink.value() = sum_weights;
-  }
-  if (exports.max_dist.valid()) {
-    assign_pos_of(dwi, 0, 3).to(exports.max_dist);
-    exports.max_dist.value() = neighbourhood.max_distance;
-  }
-  if (exports.voxelcount.valid()) {
-    assign_pos_of(dwi, 0, 3).to(exports.voxelcount);
-    exports.voxelcount.value() = n;
-  }
-}
-
-template <typename F> void Functor<F>::load_data(Image<F> &image, const std::vector<Kernel::Voxel> &voxels) {
-  const Kernel::Voxel::index_type pos({image.index(0), image.index(1), image.index(2)});
-  for (ssize_t i = 0; i != voxels.size(); ++i) {
-    assign_pos_of(voxels[i].index, 0, 3).to(image);
-    X.col(i) = image.row(3);
-  }
-  assign_pos_of(pos, 0, 3).to(image);
-}
-
-} // namespace MR::Denoise
diff --git a/src/denoise/recon.cpp b/src/denoise/recon.cpp
new file mode 100644
index 0000000000..4304e31aa8
--- /dev/null
+++ b/src/denoise/recon.cpp
@@ -0,0 +1,163 @@
+/* Copyright (c) 2008-2024 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#include "denoise/recon.h"
+
+#include "math/math.h"
+
+namespace MR::Denoise {
+
+template <typename F>
+Recon<F>::Recon(const Header &header,
+                Image<bool> &mask,
+                std::shared_ptr<Kernel::Base> kernel,
+                std::shared_ptr<Estimator::Base> estimator,
+                filter_type filter,
+                aggregator_type aggregator,
+                Exports &exports)
+    : Estimate<F>(header, mask, kernel, estimator, exports),
+      filter(filter),
+      aggregator(aggregator),
+      // FWHM = 2 x cube root of voxel spacings
+      gaussian_multiplier(-std::log(2.0) /
+                          Math::pow2(std::cbrt(header.spacing(0) * header.spacing(1) * header.spacing(2)))),
+      w(std::min(Estimate<F>::m, kernel->estimated_size())) {}
+
+template <typename F> void Recon<F>::operator()(Image<F> &dwi, Image<F> &out) {
+
+  Estimate<F> (*this)(dwi);
+
+  const ssize_t n = Estimate<F>::neighbourhood.voxels.size();
+  const ssize_t r = std::min(Estimate<F>::m, n);
+  const ssize_t q = std::max(Estimate<F>::m, n);
+  const ssize_t in_rank = r - Estimate<F>::threshold.cutoff_p;
+
+  if (r > w.size())
+    w.resize(r);
+#ifndef NDEBUG
+  w.fill(std::numeric_limits<default_type>::signaling_NaN());
+#endif
+
+  // Generate weights vector
+  double sum_weights = 0.0;
+  ssize_t out_rank = 0;
+  switch (filter) {
+  case filter_type::TRUNCATE:
+    out_rank = in_rank;
+    w.head(Estimate<F>::threshold.cutoff_p).setZero();
+    w.segment(Estimate<F>::threshold.cutoff_p, in_rank).setOnes();
+    sum_weights = double(out_rank);
+    break;
+  case filter_type::FROBENIUS: {
+    const double beta = r / q;
+    const double transition = 1.0 + std::sqrt(beta);
+    double clam = 0.0;
+    for (ssize_t i = 0; i != r; ++i) {
+      const double lam = std::max(Estimate<F>::s[i], 0.0) / q;
+      clam += lam;
+      const double y = clam / (Estimate<F>::threshold.sigma2 * (i + 1));
+      double nu = 0.0;
+      if (y > transition) {
+        nu = std::sqrt(Math::pow2(Math::pow2(y) - beta - 1.0) - (4.0 * beta)) / y;
+        ++out_rank;
+      }
+      w[i] = nu / y;
+      sum_weights += w[i];
+    }
+  } break;
+  default:
+    assert(false);
+  }
+
+  // recombine data using only eigenvectors above threshold
+  // If only the data computed when this voxel was the centre of the patch
+  //   is to be used for synthesis of the output image,
+  //   then only that individual column needs to be reconstructed;
+  //   if however the result from this patch is to contribute to the synthesized image
+  //   for all voxels that were utilised within this patch,
+  //   then we need to instead compute the full projection
+  // TODO Use a new data member local to Recon<>
+  switch (aggregator) {
+  case aggregator_type::EXCLUSIVE:
+    if (Estimate<F>::m <= n)
+      Estimate<F>::X.col(Estimate<F>::neighbourhood.centre_index) =
+          Estimate<F>::eig.eigenvectors() *
+          (w.head(r).cast<F>().matrix().asDiagonal() *
+           (Estimate<F>::eig.eigenvectors().adjoint() * Estimate<F>::X.col(Estimate<F>::neighbourhood.centre_index)));
+    else
+      Estimate<F>::X.col(Estimate<F>::neighbourhood.centre_index) =
+          Estimate<F>::X.leftCols(n) *
+          (Estimate<F>::eig.eigenvectors() *
+           (w.head(r).cast<F>().matrix().asDiagonal() *
+            Estimate<F>::eig.eigenvectors().adjoint().col(Estimate<F>::neighbourhood.centre_index)));
+    assign_pos_of(dwi).to(out);
+    out.row(3) = Estimate<F>::X.col(Estimate<F>::neighbourhood.centre_index);
+    if (Estimate<F>::exports.sum_aggregation.valid()) {
+      assign_pos_of(dwi, 0, 3).to(Estimate<F>::exports.sum_aggregation);
+      Estimate<F>::exports.sum_aggregation.value() = 1.0;
+    }
+    if (Estimate<F>::exports.rank_output.valid()) {
+      assign_pos_of(dwi, 0, 3).to(Estimate<F>::exports.rank_output);
+      Estimate<F>::exports.rank_output.value() = out_rank;
+    }
+    break;
+  default: {
+    if (Estimate<F>::m <= n)
+      Estimate<F>::X = Estimate<F>::eig.eigenvectors() * (w.head(r).cast<F>().matrix().asDiagonal() *
+                                                          (Estimate<F>::eig.eigenvectors().adjoint() * Estimate<F>::X));
+    else
+      Estimate<F>::X.leftCols(n) =
+          Estimate<F>::X.leftCols(n) * (Estimate<F>::eig.eigenvectors() * (w.head(r).cast<F>().matrix().asDiagonal() *
+                                                                           Estimate<F>::eig.eigenvectors().adjoint()));
+    std::lock_guard<std::mutex> lock(mutex_aggregator);
+    for (size_t voxel_index = 0; voxel_index != Estimate<F>::neighbourhood.voxels.size(); ++voxel_index) {
+      assign_pos_of(Estimate<F>::neighbourhood.voxels[voxel_index].index, 0, 3).to(out);
+      assign_pos_of(Estimate<F>::neighbourhood.voxels[voxel_index].index).to(Estimate<F>::exports.sum_aggregation);
+      double weight = std::numeric_limits<double>::signaling_NaN();
+      switch (aggregator) {
+      case aggregator_type::EXCLUSIVE:
+        assert(false);
+        break;
+      case aggregator_type::GAUSSIAN:
+        weight = std::exp(gaussian_multiplier * Estimate<F>::neighbourhood.voxels[voxel_index].sq_distance);
+        break;
+      case aggregator_type::INVL0:
+        weight = 1.0 / (1 + out_rank);
+        break;
+      case aggregator_type::RANK:
+        weight = out_rank;
+        break;
+      case aggregator_type::UNIFORM:
+        weight = 1.0;
+        break;
+      }
+      out.row(3) += weight * Estimate<F>::X.col(voxel_index);
+      Estimate<F>::exports.sum_aggregation.value() += weight;
+      if (Estimate<F>::exports.rank_output.valid()) {
+        assign_pos_of(Estimate<F>::neighbourhood.voxels[voxel_index].index, 0, 3).to(Estimate<F>::exports.rank_output);
+        Estimate<F>::exports.rank_output.value() += weight * out_rank;
+      }
+    }
+  } break;
+  }
+
+  if (Estimate<F>::exports.sum_optshrink.valid()) {
+    assign_pos_of(dwi, 0, 3).to(Estimate<F>::exports.sum_optshrink);
+    Estimate<F>::exports.sum_optshrink.value() = sum_weights;
+  }
+}
+
+} // namespace MR::Denoise
diff --git a/src/denoise/recon.h b/src/denoise/recon.h
new file mode 100644
index 0000000000..946d01d711
--- /dev/null
+++ b/src/denoise/recon.h
@@ -0,0 +1,67 @@
+/* Copyright (c) 2008-2024 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#pragma once
+
+#include <limits>
+#include <memory>
+#include <mutex>
+
+#include <Eigen/Dense>
+
+#include "denoise/estimate.h"
+#include "denoise/estimator/base.h"
+#include "denoise/exports.h"
+#include "denoise/kernel/base.h"
+#include "header.h"
+#include "image.h"
+
+namespace MR::Denoise {
+
+template <typename F> class Recon : public Estimate<F> {
+
+public:
+  Recon(const Header &header,
+        Image<bool> &mask,
+        std::shared_ptr<Kernel::Base> kernel,
+        std::shared_ptr<Estimator::Base> estimator,
+        filter_type filter,
+        aggregator_type aggregator,
+        Exports &exports);
+
+  void operator()(Image<F> &dwi, Image<F> &out);
+
+protected:
+  // Denoising configuration
+  filter_type filter;
+  aggregator_type aggregator;
+  double gaussian_multiplier;
+
+  // Reusable memory
+  vector_type w;
+
+  // Some data can only be written in a thread-safe manner
+  static std::mutex mutex_aggregator;
+};
+
+template <typename F> std::mutex Recon<F>::mutex_aggregator;
+
+template class Recon<float>;
+template class Recon<cfloat>;
+template class Recon<double>;
+template class Recon<cdouble>;
+
+} // namespace MR::Denoise

From 5d19ae0c4db3267c993d45a5e034f1c7ea7bf859 Mon Sep 17 00:00:00 2001
From: Robert Smith <robert.smith@florey.edu.au>
Date: Sun, 17 Nov 2024 22:37:09 +1100
Subject: [PATCH 19/34] dwidenoise: Separate members for input vs denoised data

---
 src/denoise/recon.cpp | 45 ++++++++++++++++++++++++++-----------------
 src/denoise/recon.h   |  1 +
 2 files changed, 28 insertions(+), 18 deletions(-)

diff --git a/src/denoise/recon.cpp b/src/denoise/recon.cpp
index 4304e31aa8..19ebdaf19b 100644
--- a/src/denoise/recon.cpp
+++ b/src/denoise/recon.cpp
@@ -34,11 +34,12 @@ Recon<F>::Recon(const Header &header,
       // FWHM = 2 x cube root of voxel spacings
       gaussian_multiplier(-std::log(2.0) /
                           Math::pow2(std::cbrt(header.spacing(0) * header.spacing(1) * header.spacing(2)))),
-      w(std::min(Estimate<F>::m, kernel->estimated_size())) {}
+      w(std::min(Estimate<F>::m, kernel->estimated_size())),
+      Xr(Estimate<F>::m, aggregator == aggregator_type::EXCLUSIVE ? 1 : kernel->estimated_size()) {}
 
 template <typename F> void Recon<F>::operator()(Image<F> &dwi, Image<F> &out) {
 
-  Estimate<F> (*this)(dwi);
+  Estimate<F>::operator()(dwi);
 
   const ssize_t n = Estimate<F>::neighbourhood.voxels.size();
   const ssize_t r = std::min(Estimate<F>::m, n);
@@ -47,8 +48,11 @@ template <typename F> void Recon<F>::operator()(Image<F> &dwi, Image<F> &out) {
 
   if (r > w.size())
     w.resize(r);
+  if (aggregator != aggregator_type::EXCLUSIVE && n > Xr.cols())
+    Xr.resize(Estimate<F>::m, n);
 #ifndef NDEBUG
   w.fill(std::numeric_limits<default_type>::signaling_NaN());
+  Xr.fill(std::numeric_limits<default_type>::signaling_NaN());
 #endif
 
   // Generate weights vector
@@ -93,18 +97,19 @@ template <typename F> void Recon<F>::operator()(Image<F> &dwi, Image<F> &out) {
   switch (aggregator) {
   case aggregator_type::EXCLUSIVE:
     if (Estimate<F>::m <= n)
-      Estimate<F>::X.col(Estimate<F>::neighbourhood.centre_index) =
-          Estimate<F>::eig.eigenvectors() *
-          (w.head(r).cast<F>().matrix().asDiagonal() *
-           (Estimate<F>::eig.eigenvectors().adjoint() * Estimate<F>::X.col(Estimate<F>::neighbourhood.centre_index)));
+      Xr.noalias() =                                                       //
+          Estimate<F>::eig.eigenvectors() *                                //
+          (w.head(r).cast<F>().matrix().asDiagonal() *                     //
+           (Estimate<F>::eig.eigenvectors().adjoint() *                    //
+            Estimate<F>::X.col(Estimate<F>::neighbourhood.centre_index))); //
     else
-      Estimate<F>::X.col(Estimate<F>::neighbourhood.centre_index) =
-          Estimate<F>::X.leftCols(n) *
-          (Estimate<F>::eig.eigenvectors() *
-           (w.head(r).cast<F>().matrix().asDiagonal() *
-            Estimate<F>::eig.eigenvectors().adjoint().col(Estimate<F>::neighbourhood.centre_index)));
+      Xr.noalias() =                                                                                  //
+          Estimate<F>::X.leftCols(n) *                                                                //
+          (Estimate<F>::eig.eigenvectors() *                                                          //
+           (w.head(r).cast<F>().matrix().asDiagonal() *                                               //
+            Estimate<F>::eig.eigenvectors().adjoint().col(Estimate<F>::neighbourhood.centre_index))); //
     assign_pos_of(dwi).to(out);
-    out.row(3) = Estimate<F>::X.col(Estimate<F>::neighbourhood.centre_index);
+    out.row(3) = Xr.col(0);
     if (Estimate<F>::exports.sum_aggregation.valid()) {
       assign_pos_of(dwi, 0, 3).to(Estimate<F>::exports.sum_aggregation);
       Estimate<F>::exports.sum_aggregation.value() = 1.0;
@@ -116,12 +121,16 @@ template <typename F> void Recon<F>::operator()(Image<F> &dwi, Image<F> &out) {
     break;
   default: {
     if (Estimate<F>::m <= n)
-      Estimate<F>::X = Estimate<F>::eig.eigenvectors() * (w.head(r).cast<F>().matrix().asDiagonal() *
-                                                          (Estimate<F>::eig.eigenvectors().adjoint() * Estimate<F>::X));
+      Xr.noalias() =                                                      //
+          Estimate<F>::eig.eigenvectors() *                               //
+          (w.head(r).cast<F>().matrix().asDiagonal() *                    //
+           (Estimate<F>::eig.eigenvectors().adjoint() * Estimate<F>::X)); //
     else
-      Estimate<F>::X.leftCols(n) =
-          Estimate<F>::X.leftCols(n) * (Estimate<F>::eig.eigenvectors() * (w.head(r).cast<F>().matrix().asDiagonal() *
-                                                                           Estimate<F>::eig.eigenvectors().adjoint()));
+      Xr.leftCols(n).noalias() =                         //
+          Estimate<F>::X.leftCols(n) *                   //
+          (Estimate<F>::eig.eigenvectors() *             //
+           (w.head(r).cast<F>().matrix().asDiagonal() *  //
+            Estimate<F>::eig.eigenvectors().adjoint())); //
     std::lock_guard<std::mutex> lock(mutex_aggregator);
     for (size_t voxel_index = 0; voxel_index != Estimate<F>::neighbourhood.voxels.size(); ++voxel_index) {
       assign_pos_of(Estimate<F>::neighbourhood.voxels[voxel_index].index, 0, 3).to(out);
@@ -144,7 +153,7 @@ template <typename F> void Recon<F>::operator()(Image<F> &dwi, Image<F> &out) {
         weight = 1.0;
         break;
       }
-      out.row(3) += weight * Estimate<F>::X.col(voxel_index);
+      out.row(3) += weight * Xr.col(voxel_index);
       Estimate<F>::exports.sum_aggregation.value() += weight;
       if (Estimate<F>::exports.rank_output.valid()) {
         assign_pos_of(Estimate<F>::neighbourhood.voxels[voxel_index].index, 0, 3).to(Estimate<F>::exports.rank_output);
diff --git a/src/denoise/recon.h b/src/denoise/recon.h
index 946d01d711..8646d5b56d 100644
--- a/src/denoise/recon.h
+++ b/src/denoise/recon.h
@@ -52,6 +52,7 @@ template <typename F> class Recon : public Estimate<F> {
 
   // Reusable memory
   vector_type w;
+  typename Estimate<F>::MatrixType Xr;
 
   // Some data can only be written in a thread-safe manner
   static std::mutex mutex_aggregator;

From e6c81f3090d85149b3b7bc7ba0e2928160841c1a Mon Sep 17 00:00:00 2001
From: Robert Smith <robert.smith@florey.edu.au>
Date: Thu, 21 Nov 2024 21:49:55 +1100
Subject: [PATCH 20/34] dwidenoise & dwi2noise: Add subsampling capability

Closes #3034.
---
 cmd/dwi2noise.cpp                  |  30 +++++---
 cmd/dwidenoise.cpp                 |  39 ++++++++---
 src/denoise/estimate.cpp           |  29 ++++++--
 src/denoise/estimate.h             |   9 +++
 src/denoise/exports.h              |  44 ++++++++----
 src/denoise/kernel/cuboid.cpp      |  57 +++++++++------
 src/denoise/kernel/cuboid.h        |   7 +-
 src/denoise/kernel/kernel.cpp      |  77 +++++++++++++-------
 src/denoise/kernel/kernel.h        |   7 +-
 src/denoise/kernel/sphere_base.cpp |  55 ++++++++++++---
 src/denoise/kernel/sphere_base.h   |   7 +-
 src/denoise/kernel/sphere_radius.h |   6 +-
 src/denoise/kernel/sphere_ratio.h  |   4 +-
 src/denoise/recon.cpp              |  72 ++++++++++++-------
 src/denoise/recon.h                |   7 +-
 src/denoise/subsample.cpp          | 108 +++++++++++++++++++++++++++++
 src/denoise/subsample.h            |  60 ++++++++++++++++
 17 files changed, 482 insertions(+), 136 deletions(-)
 create mode 100644 src/denoise/subsample.cpp
 create mode 100644 src/denoise/subsample.h

diff --git a/cmd/dwi2noise.cpp b/cmd/dwi2noise.cpp
index 045dd5dd2a..eda4d4fde6 100644
--- a/cmd/dwi2noise.cpp
+++ b/cmd/dwi2noise.cpp
@@ -22,6 +22,7 @@
 #include "denoise/estimator/estimator.h"
 #include "denoise/exports.h"
 #include "denoise/kernel/kernel.h"
+#include "denoise/subsample.h"
 #include "exception.h"
 
 using namespace MR;
@@ -60,7 +61,9 @@ void usage() {
 
   + Kernel::shape_description
 
-  + Kernel::size_description;
+  + Kernel::default_size_description
+
+  + Kernel::cuboid_size_description;
 
   AUTHOR = "Daan Christiaens (daan.christiaens@kcl.ac.uk)"
            " and Jelle Veraart (jelle.veraart@nyumc.org)"
@@ -90,6 +93,7 @@ void usage() {
   + datatype_option
   + Estimator::option
   + Kernel::options
+  + subsample_option
 
   // TODO Implement mask option
   // Note that behaviour of -mask for dwi2noise may be different to that of dwidenoise
@@ -104,6 +108,9 @@ void usage() {
     + Argument("image").type_image_out()
   + Option("voxelcount",
            "The number of voxels that contributed to the PCA for processing of each voxel")
+    + Argument("image").type_image_out()
+  + Option("patchcount",
+           "The number of unique patches to which an image voxel contributes")
     + Argument("image").type_image_out();
 
   COPYRIGHT =
@@ -137,12 +144,13 @@ void usage() {
 
 template <typename T>
 void run(Header &data,
+         std::shared_ptr<Subsample> subsample,
          std::shared_ptr<Kernel::Base> kernel,
          std::shared_ptr<Estimator::Base> estimator,
          Exports &exports) {
   auto input = data.get_image<T>().with_direct_io(3);
   Image<bool> mask; // unused
-  Estimate<T> func(data, mask, kernel, estimator, exports);
+  Estimate<T> func(data, mask, subsample, kernel, estimator, exports);
   ThreadedLoop("running MP-PCA noise level estimation", data, 0, 3).run(func, input);
 }
 
@@ -152,13 +160,16 @@ void run() {
   if (dwi.ndim() != 4 || dwi.size(3) <= 1)
     throw Exception("input image must be 4-dimensional");
 
-  auto kernel = Kernel::make_kernel(dwi);
+  auto subsample = Subsample::make(dwi);
+  assert(subsample);
+
+  auto kernel = Kernel::make_kernel(dwi, subsample->get_factors());
   assert(kernel);
 
   auto estimator = Estimator::make_estimator();
   assert(estimator);
 
-  Exports exports(dwi);
+  Exports exports(dwi, subsample->header());
   exports.set_noise_out(argument[1]);
   auto opt = get_options("rank");
   if (!opt.empty())
@@ -169,6 +180,9 @@ void run() {
   opt = get_options("voxelcount");
   if (!opt.empty())
     exports.set_voxelcount(opt[0][0]);
+  opt = get_options("patchcount");
+  if (!opt.empty())
+    exports.set_patchcount(opt[0][0]);
 
   int prec = get_option_value("datatype", 0); // default: single precision
   if (dwi.datatype().is_complex())
@@ -176,19 +190,19 @@ void run() {
   switch (prec) {
   case 0:
     INFO("select real float32 for processing");
-    run<float>(dwi, kernel, estimator, exports);
+    run<float>(dwi, subsample, kernel, estimator, exports);
     break;
   case 1:
     INFO("select real float64 for processing");
-    run<double>(dwi, kernel, estimator, exports);
+    run<double>(dwi, subsample, kernel, estimator, exports);
     break;
   case 2:
     INFO("select complex float32 for processing");
-    run<cfloat>(dwi, kernel, estimator, exports);
+    run<cfloat>(dwi, subsample, kernel, estimator, exports);
     break;
   case 3:
     INFO("select complex float64 for processing");
-    run<cdouble>(dwi, kernel, estimator, exports);
+    run<cdouble>(dwi, subsample, kernel, estimator, exports);
     break;
   }
 }
diff --git a/cmd/dwidenoise.cpp b/cmd/dwidenoise.cpp
index ad55fc03a9..3a38e9591c 100644
--- a/cmd/dwidenoise.cpp
+++ b/cmd/dwidenoise.cpp
@@ -15,7 +15,6 @@
  */
 
 #include <string>
-#include <vector>
 
 #include "command.h"
 #include "header.h"
@@ -37,6 +36,7 @@
 #include "denoise/kernel/sphere_radius.h"
 #include "denoise/kernel/sphere_ratio.h"
 #include "denoise/recon.h"
+#include "denoise/subsample.h"
 
 using namespace MR;
 using namespace App;
@@ -67,7 +67,9 @@ void usage() {
 
   + Kernel::shape_description
 
-  + Kernel::size_description
+  + Kernel::default_size_description
+
+  + Kernel::cuboid_size_description
 
   + "By default, optimal value shrinkage based on minimisation of the Frobenius norm "
     "will be used to attenuate eigenvectors based on the estimated noise level. "
@@ -130,6 +132,7 @@ void usage() {
   + datatype_option
   + Estimator::option
   + Kernel::options
+  + subsample_option
 
   + OptionGroup("Options that affect reconstruction of the output image series")
   // TODO Separate masks for voxels to contribute to patches vs. voxels for which to perform denoising
@@ -178,7 +181,7 @@ void usage() {
            " so a scale factor sqrt(2) applies.")
     + Argument("image").type_image_out()
   + Option("rank_input",
-           "The signal rank estimated for the denoising patch centred at each input image voxel")
+           "The signal rank estimated for each denoising patch")
     + Argument("image").type_image_out()
   + Option("rank_output",
            "An estimated rank for the output image data, accounting for multi-patch aggregation")
@@ -191,6 +194,9 @@ void usage() {
   + Option("voxelcount",
            "The number of voxels that contributed to the PCA for processing of each voxel")
     + Argument("image").type_image_out()
+  + Option("patchcount",
+           "The number of unique patches to which an image voxel contributes")
+    + Argument("image").type_image_out()
   + Option("sum_aggregation",
            "The sum of aggregation weights of those patches contributing to each output voxel")
     + Argument("image").type_image_out()
@@ -236,6 +242,7 @@ std::complex<double> operator/(const std::complex<double> &c, const float n) { r
 template <typename T>
 void run(Header &data,
          Image<bool> &mask,
+         std::shared_ptr<Subsample> subsample,
          std::shared_ptr<Kernel::Base> kernel,
          std::shared_ptr<Estimator::Base> estimator,
          filter_type filter,
@@ -248,7 +255,7 @@ void run(Header &data,
   header.datatype() = DataType::from<T>();
   auto output = Image<T>::create(output_name, header);
   // run
-  Recon<T> func(data, mask, kernel, estimator, filter, aggregator, exports);
+  Recon<T> func(data, mask, subsample, kernel, estimator, filter, aggregator, exports);
   ThreadedLoop("running MP-PCA denoising", data, 0, 3).run(func, input, output);
   // Rescale output if performing aggregation
   if (aggregator == aggregator_type::EXCLUSIVE)
@@ -276,7 +283,10 @@ void run() {
     check_dimensions(mask, dwi, 0, 3);
   }
 
-  auto kernel = Kernel::make_kernel(dwi);
+  auto subsample = Subsample::make(dwi);
+  assert(subsample);
+
+  auto kernel = Kernel::make_kernel(dwi, subsample->get_factors());
   assert(kernel);
 
   auto estimator = Estimator::make_estimator();
@@ -289,10 +299,14 @@ void run() {
 
   aggregator_type aggregator = aggregator_type::GAUSSIAN;
   opt = get_options("aggregator");
-  if (!opt.empty())
+  if (!opt.empty()) {
     aggregator = aggregator_type(int(opt[0][0]));
+    if (aggregator == aggregator_type::EXCLUSIVE && subsample->get_factors() != std::array<ssize_t, 3>({1, 1, 1}))
+      throw Exception("Cannot combine -aggregator exclusive with subsampling; "
+                      "would result in empty output voxels");
+  }
 
-  Exports exports(dwi);
+  Exports exports(dwi, subsample->header());
   opt = get_options("noise_out");
   if (!opt.empty())
     exports.set_noise_out(opt[0][0]);
@@ -323,6 +337,9 @@ void run() {
   opt = get_options("voxelcount");
   if (!opt.empty())
     exports.set_voxelcount(opt[0][0]);
+  opt = get_options("patchcount");
+  if (!opt.empty())
+    exports.set_patchcount(opt[0][0]);
 
   opt = get_options("sum_aggregation");
   if (!opt.empty()) {
@@ -341,19 +358,19 @@ void run() {
   switch (prec) {
   case 0:
     INFO("select real float32 for processing");
-    run<float>(dwi, mask, kernel, estimator, filter, aggregator, argument[1], exports);
+    run<float>(dwi, mask, subsample, kernel, estimator, filter, aggregator, argument[1], exports);
     break;
   case 1:
     INFO("select real float64 for processing");
-    run<double>(dwi, mask, kernel, estimator, filter, aggregator, argument[1], exports);
+    run<double>(dwi, mask, subsample, kernel, estimator, filter, aggregator, argument[1], exports);
     break;
   case 2:
     INFO("select complex float32 for processing");
-    run<cfloat>(dwi, mask, kernel, estimator, filter, aggregator, argument[1], exports);
+    run<cfloat>(dwi, mask, subsample, kernel, estimator, filter, aggregator, argument[1], exports);
     break;
   case 3:
     INFO("select complex float64 for processing");
-    run<cdouble>(dwi, mask, kernel, estimator, filter, aggregator, argument[1], exports);
+    run<cdouble>(dwi, mask, subsample, kernel, estimator, filter, aggregator, argument[1], exports);
     break;
   }
 }
diff --git a/src/denoise/estimate.cpp b/src/denoise/estimate.cpp
index c2c19a269a..a7306705e5 100644
--- a/src/denoise/estimate.cpp
+++ b/src/denoise/estimate.cpp
@@ -25,11 +25,13 @@ namespace MR::Denoise {
 template <typename F>
 Estimate<F>::Estimate(const Header &header,
                       Image<bool> &mask,
+                      std::shared_ptr<Subsample> subsample,
                       std::shared_ptr<Kernel::Base> kernel,
                       std::shared_ptr<Estimator::Base> estimator,
                       Exports &exports)
     : m(header.size(3)),
       mask(mask),
+      subsample(subsample),
       kernel(kernel),
       estimator(estimator),
       X(m, kernel->estimated_size()),
@@ -40,6 +42,17 @@ Estimate<F>::Estimate(const Header &header,
 
 template <typename F> void Estimate<F>::operator()(Image<F> &dwi) {
 
+  // There are two options here for looping in the presence of subsampling:
+  // 1. Loop over the input image
+  //    Skip voxels that don't lie at the centre of a patch
+  //    Have to transform input image voxel indices to subsampled image voxel indices for some optional outputs
+  // 2. Loop over the subsampled image
+  //    In some use cases there may not be any image created that conforms to this voxel grid
+  //    Have to transform the subsampled voxel index into an input image voxel index for the centre of the patch
+  // Going to go with 1. for now, as for 2. may not have a suitable image over which to loop
+  if (!subsample->process(Kernel::Voxel::index_type({dwi.index(0), dwi.index(1), dwi.index(2)})))
+    return;
+
   // Process voxels in mask only
   if (mask.valid()) {
     assign_pos_of(dwi, 0, 3).to(mask);
@@ -90,22 +103,30 @@ template <typename F> void Estimate<F>::operator()(Image<F> &dwi) {
   const ssize_t in_rank = r - threshold.cutoff_p;
 
   // Store additional output maps if requested
+  auto ss_index = subsample->in2ss({dwi.index(0), dwi.index(1), dwi.index(2)});
   if (exports.noise_out.valid()) {
-    assign_pos_of(dwi, 0, 3).to(exports.noise_out);
+    assign_pos_of(ss_index).to(exports.noise_out);
     exports.noise_out.value() = float(std::sqrt(threshold.sigma2));
   }
   if (exports.rank_input.valid()) {
-    assign_pos_of(dwi, 0, 3).to(exports.rank_input);
+    assign_pos_of(ss_index).to(exports.rank_input);
     exports.rank_input.value() = in_rank;
   }
   if (exports.max_dist.valid()) {
-    assign_pos_of(dwi, 0, 3).to(exports.max_dist);
+    assign_pos_of(ss_index).to(exports.max_dist);
     exports.max_dist.value() = neighbourhood.max_distance;
   }
   if (exports.voxelcount.valid()) {
-    assign_pos_of(dwi, 0, 3).to(exports.voxelcount);
+    assign_pos_of(ss_index).to(exports.voxelcount);
     exports.voxelcount.value() = n;
   }
+  if (exports.patchcount.valid()) {
+    std::lock_guard<std::mutex> lock(Estimate<F>::mutex);
+    for (const auto &v : neighbourhood.voxels) {
+      assign_pos_of(v.index).to(exports.patchcount);
+      exports.patchcount.value() = exports.patchcount.value() + 1;
+    }
+  }
 }
 
 template <typename F> void Estimate<F>::load_data(Image<F> &image, const std::vector<Kernel::Voxel> &voxels) {
diff --git a/src/denoise/estimate.h b/src/denoise/estimate.h
index babbabfb60..c8f26f423c 100644
--- a/src/denoise/estimate.h
+++ b/src/denoise/estimate.h
@@ -17,6 +17,7 @@
 #pragma once
 
 #include <memory>
+#include <mutex>
 
 #include <Eigen/Dense>
 
@@ -27,6 +28,7 @@
 #include "denoise/kernel/base.h"
 #include "denoise/kernel/data.h"
 #include "denoise/kernel/voxel.h"
+#include "denoise/subsample.h"
 #include "header.h"
 #include "image.h"
 
@@ -39,6 +41,7 @@ template <typename F> class Estimate {
 
   Estimate(const Header &header,
            Image<bool> &mask,
+           std::shared_ptr<Subsample> subsample,
            std::shared_ptr<Kernel::Base> kernel,
            std::shared_ptr<Estimator::Base> estimator,
            Exports &exports);
@@ -50,6 +53,7 @@ template <typename F> class Estimate {
 
   // Denoising configuration
   Image<bool> mask;
+  std::shared_ptr<Subsample> subsample;
   std::shared_ptr<Kernel::Base> kernel;
   std::shared_ptr<Estimator::Base> estimator;
 
@@ -65,9 +69,14 @@ template <typename F> class Estimate {
   // Export images
   Exports exports;
 
+  // Some data can only be written in a thread-safe manner
+  static std::mutex mutex;
+
   void load_data(Image<F> &image, const std::vector<Kernel::Voxel> &voxels);
 };
 
+template <typename F> std::mutex Estimate<F>::mutex;
+
 template class Estimate<float>;
 template class Estimate<cfloat>;
 template class Estimate<double>;
diff --git a/src/denoise/exports.h b/src/denoise/exports.h
index 72ee180e2a..d93008025c 100644
--- a/src/denoise/exports.h
+++ b/src/denoise/exports.h
@@ -25,21 +25,39 @@ namespace MR::Denoise {
 
 class Exports {
 public:
-  Exports(const Header &in) : H(in) {
-    H.ndim() = 3;
-    H.reset_intensity_scaling();
+  Exports(const Header &in, const Header &ss) : H_in(in), H_ss(ss) {
+    H_in.ndim() = 3;
+    H_in.reset_intensity_scaling();
+    H_in.datatype() = DataType::Float32;
+    H_in.datatype().set_byte_order_native();
+  }
+  void set_noise_out(const std::string &path) { noise_out = Image<float>::create(path, H_ss); }
+  void set_rank_input(const std::string &path) {
+    Header H(H_ss);
+    H.datatype() = DataType::UInt16;
+    H.datatype().set_byte_order_native();
+    rank_input = Image<uint16_t>::create(path, H);
+  }
+  void set_rank_output(const std::string &path) { rank_output = Image<float>::create(path, H_in); }
+  void set_sum_optshrink(const std::string &path) { sum_optshrink = Image<float>::create(path, H_ss); }
+  void set_max_dist(const std::string &path) { max_dist = Image<float>::create(path, H_ss); }
+  void set_voxelcount(const std::string &path) {
+    Header H(H_ss);
+    H.datatype() = DataType::UInt16;
+    H.datatype().set_byte_order_native();
+    voxelcount = Image<uint16_t>::create(path, H);
+  }
+  void set_patchcount(const std::string &path) {
+    Header H(H_in);
+    H.datatype() = DataType::UInt16;
+    H.datatype().set_byte_order_native();
+    patchcount = Image<uint16_t>::create(path, H);
   }
-  void set_noise_out(const std::string &path) { noise_out = Image<float>::create(path, H); }
-  void set_rank_input(const std::string &path) { rank_input = Image<uint16_t>::create(path, H); }
-  void set_rank_output(const std::string &path) { rank_output = Image<float>::create(path, H); }
-  void set_sum_optshrink(const std::string &path) { sum_optshrink = Image<float>::create(path, H); }
-  void set_max_dist(const std::string &path) { max_dist = Image<float>::create(path, H); }
-  void set_voxelcount(const std::string &path) { voxelcount = Image<uint16_t>::create(path, H); }
   void set_sum_aggregation(const std::string &path) {
     if (path.empty())
-      sum_aggregation = Image<float>::scratch(H, "Scratch image for patch aggregation sums");
+      sum_aggregation = Image<float>::scratch(H_in, "Scratch image for patch aggregation sums");
     else
-      sum_aggregation = Image<float>::create(path, H);
+      sum_aggregation = Image<float>::create(path, H_in);
   }
 
   Image<float> noise_out;
@@ -48,10 +66,12 @@ class Exports {
   Image<float> sum_optshrink;
   Image<float> max_dist;
   Image<uint16_t> voxelcount;
+  Image<uint16_t> patchcount;
   Image<float> sum_aggregation;
 
 protected:
-  Header H;
+  Header H_in;
+  Header H_ss;
 };
 
 } // namespace MR::Denoise
diff --git a/src/denoise/kernel/cuboid.cpp b/src/denoise/kernel/cuboid.cpp
index 261c97fbfa..b65b579a3a 100644
--- a/src/denoise/kernel/cuboid.cpp
+++ b/src/denoise/kernel/cuboid.cpp
@@ -18,25 +18,42 @@
 
 namespace MR::Denoise::Kernel {
 
-Cuboid::Cuboid(const Header &header, const std::vector<uint32_t> &extent)
+Cuboid::Cuboid(const Header &header,
+               const std::array<ssize_t, 3> &extent,
+               const std::array<ssize_t, 3> &subsample_factors)
     : Base(header),
-      half_extent({ssize_t(extent[0] / 2), ssize_t(extent[1] / 2), ssize_t(extent[2] / 2)}),
-      size(ssize_t(extent[0]) * ssize_t(extent[1]) * ssize_t(extent[2])),
-      centre_index(size / 2) {
-  for (auto e : extent) {
-    if (!(e % 2))
-      throw Exception("Size of cubic kernel must be an odd integer");
+      size(extent[0] * extent[1] * extent[2]),
+      // Only sensible if no subsampling is performed,
+      //   and every single DWI voxel is reconstructed from a patch centred at that voxel,
+      //   with no overcomplete local PCA (aggregator == EXCLUSIVE)
+      centre_index(subsample_factors == std::array<ssize_t, 3>({1, 1, 1}) ? (size / 2) : -1) {
+  for (ssize_t axis = 0; axis != 3; ++axis) {
+    if (subsample_factors[axis] % 2) {
+      if (!(extent[axis] % 2))
+        throw Exception("For no subsampling / subsampling by an odd number, "
+                        "size of cubic kernel must be an odd integer");
+      bounding_box(axis, 0) = -extent[axis] / 2;
+      bounding_box(axis, 1) = extent[axis] / 2;
+      halfvoxel_offsets[axis] = 0.0;
+    } else {
+      if (extent[axis] % 2)
+        throw Exception("For subsampling by an even number, "
+                        "size of cubic kernel must be an even integer");
+      bounding_box(axis, 0) = 1 - extent[axis] / 2;
+      bounding_box(axis, 1) = extent[axis] / 2;
+      halfvoxel_offsets[axis] = 0.5;
+    }
   }
 }
 
 namespace {
 // patch handling at image edges
-inline ssize_t wrapindex(int p, int r, int e, int max) {
+inline ssize_t wrapindex(int p, int r, int bbminus, int bbplus, int max) {
   int rr = p + r;
   if (rr < 0)
-    rr = e - r;
+    rr = bbplus - r;
   if (rr >= max)
-    rr = (max - 1) - e - r;
+    rr = (max - 1) + bbminus - r;
   return rr;
 }
 } // namespace
@@ -45,16 +62,16 @@ Data Cuboid::operator()(const Voxel::index_type &pos) const {
   Data result(centre_index);
   Voxel::index_type voxel;
   Offset::index_type offset;
-  for (offset[2] = -half_extent[2]; offset[2] <= half_extent[2]; ++offset[2]) {
-    voxel[2] = wrapindex(pos[2], offset[2], half_extent[2], H.size(2));
-    for (offset[1] = -half_extent[1]; offset[1] <= half_extent[1]; ++offset[1]) {
-      voxel[1] = wrapindex(pos[1], offset[1], half_extent[1], H.size(1));
-      for (offset[0] = -half_extent[0]; offset[0] <= half_extent[0]; ++offset[0]) {
-        voxel[0] = wrapindex(pos[0], offset[0], half_extent[0], H.size(0));
-        // Both "pos" and "voxel" are unsigned, so beware of integer overflow
-        const default_type sq_distance = Math::pow2(std::min(pos[0] - voxel[0], voxel[0] - pos[0]) * H.spacing(0)) +
-                                         Math::pow2(std::min(pos[1] - voxel[1], voxel[1] - pos[1]) * H.spacing(1)) +
-                                         Math::pow2(std::min(pos[2] - voxel[2], voxel[2] - pos[2]) * H.spacing(2));
+  for (offset[2] = bounding_box(2, 0); offset[2] <= bounding_box(2, 1); ++offset[2]) {
+    voxel[2] = wrapindex(pos[2], offset[2], bounding_box(2, 0), bounding_box(2, 1), H.size(2));
+    for (offset[1] = bounding_box(1, 0); offset[1] <= bounding_box(1, 1); ++offset[1]) {
+      voxel[1] = wrapindex(pos[1], offset[1], bounding_box(1, 0), bounding_box(1, 1), H.size(1));
+      for (offset[0] = bounding_box(0, 0); offset[0] <= bounding_box(0, 1); ++offset[0]) {
+        voxel[0] = wrapindex(pos[0], offset[0], bounding_box(0, 0), bounding_box(0, 1), H.size(0));
+        assert(!is_out_of_bounds(H, voxel, 0, 3));
+        const default_type sq_distance = Math::pow2(pos[0] + halfvoxel_offsets[0] - voxel[0]) * H.spacing(0) +
+                                         Math::pow2(pos[1] + halfvoxel_offsets[1] - voxel[1]) * H.spacing(1) +
+                                         Math::pow2(pos[2] + halfvoxel_offsets[2] - voxel[2]) * H.spacing(2);
         result.voxels.push_back(Voxel(voxel, sq_distance));
         result.max_distance = std::max(result.max_distance, sq_distance);
       }
diff --git a/src/denoise/kernel/cuboid.h b/src/denoise/kernel/cuboid.h
index 0d2230009d..93a318e136 100644
--- a/src/denoise/kernel/cuboid.h
+++ b/src/denoise/kernel/cuboid.h
@@ -16,7 +16,7 @@
 
 #pragma once
 
-#include <vector>
+#include <array>
 
 #include "denoise/kernel/base.h"
 #include "denoise/kernel/data.h"
@@ -27,16 +27,17 @@ namespace MR::Denoise::Kernel {
 class Cuboid : public Base {
 
 public:
-  Cuboid(const Header &header, const std::vector<uint32_t> &extent);
+  Cuboid(const Header &header, const std::array<ssize_t, 3> &extent, const std::array<ssize_t, 3> &subsample_factors);
   Cuboid(const Cuboid &) = default;
   ~Cuboid() final = default;
   Data operator()(const Voxel::index_type &pos) const override;
   ssize_t estimated_size() const override { return size; }
 
 private:
-  const Voxel::index_type half_extent;
+  Eigen::Array<int, 3, 2> bounding_box;
   const ssize_t size;
   const ssize_t centre_index;
+  std::array<default_type, 3> halfvoxel_offsets;
 };
 
 } // namespace MR::Denoise::Kernel
diff --git a/src/denoise/kernel/kernel.cpp b/src/denoise/kernel/kernel.cpp
index e5cd077166..e991b86508 100644
--- a/src/denoise/kernel/kernel.cpp
+++ b/src/denoise/kernel/kernel.cpp
@@ -40,7 +40,7 @@ const char *const shape_description =
     "the centre of the kernel will be offset from the voxel being processed "
     "such that the entire volume of the kernel resides within the image FoV.";
 
-const char *const size_description =
+const char *const default_size_description =
     "The size of the default spherical kernel is set to select a number of voxels that is "
     "1.0 / 0.85 ~ 1.18 times the number of volumes in the input series. "
     "If a cuboid kernel is requested, "
@@ -50,6 +50,20 @@ const char *const size_description =
     "e.g., 5x5x5 for data with <= 125 DWI volumes, "
     "7x7x7 for data with <= 343 DWI volumes, etc.";
 
+const char *const cuboid_size_description =
+    "Permissible sizes for the cuboid kernel depend on the subsampling factor. "
+    "If no subsampling is performed, or the subsampling factor is odd, "
+    "then the extent(s) of the kernel must be odd, "
+    "such that a unique voxel lies at the very centre of each kernel. "
+    "If however an even subsampling factor is used, "
+    "then the extent(s) of the kernel must be even, "
+    "reflecting the fact that it is a voxel corner that resides at the centre of the kernel."
+    "In either case, if the extent is specified manually, "
+    "the user can either provide a single integer---"
+    "which will determine the number of voxels in the kernel across all three spatial axes---"
+    "or a comma-separated list of three integers,"
+    "individually defining the number of voxels in the kernel for all three spatial axes.";
+
 // clang-format off
 const OptionGroup options = OptionGroup("Options for controlling the sliding spatial window kernel")
 + Option("shape",
@@ -65,11 +79,11 @@ const OptionGroup options = OptionGroup("Options for controlling the sliding spa
 // TODO Command-line option that allows user to specify minimum absolute number of voxels in kernel
 + Option("extent",
          "Set the patch size of the cuboid kernel; "
-         "can be either a single odd integer or a comma-separated triplet of odd integers")
+         "can be either a single integer or a comma-separated triplet of integers (see Description)")
   + Argument("window").type_sequence_int();
 // clang-format on
 
-std::shared_ptr<Base> make_kernel(const Header &H) {
+std::shared_ptr<Base> make_kernel(const Header &H, const std::array<ssize_t, 3> &subsample_factors) {
   auto opt = App::get_options("shape");
   const Kernel::shape_type shape = opt.empty() ? Kernel::shape_type::SPHERE : Kernel::shape_type((int)(opt[0][0]));
   std::shared_ptr<Kernel::Base> kernel;
@@ -81,43 +95,54 @@ std::shared_ptr<Base> make_kernel(const Header &H) {
       throw Exception("-extent option does not apply to spherical kernel");
     opt = get_options("radius_mm");
     if (opt.empty())
-      return std::make_shared<SphereRatio>(H, get_option_value("radius_ratio", sphere_multiplier_default));
-    return std::make_shared<SphereFixedRadius>(H, opt[0][0]);
+      return std::make_shared<SphereRatio>(
+          H, get_option_value("radius_ratio", sphere_multiplier_default), subsample_factors);
+    return std::make_shared<SphereFixedRadius>(H, opt[0][0], subsample_factors);
   }
   case Kernel::shape_type::CUBOID: {
     if (!get_options("radius_mm").empty() || !get_options("radius_ratio").empty())
       throw Exception("-radius_* options are inapplicable if cuboid kernel shape is selected");
     opt = get_options("extent");
-    std::vector<uint32_t> extent;
+    std::array<ssize_t, 3> extent;
     if (!opt.empty()) {
-      extent = parse_ints<uint32_t>(opt[0][0]);
-      if (extent.size() == 1)
-        extent = {extent[0], extent[0], extent[0]};
-      if (extent.size() != 3)
+      auto userinput = parse_ints<uint32_t>(opt[0][0]);
+      if (userinput.size() == 1)
+        extent = {userinput[0], userinput[0], userinput[0]};
+      else if (userinput.size() == 3)
+        extent = {userinput[0], userinput[1], userinput[2]};
+      else
         throw Exception("-extent must be either a scalar or a list of length 3");
-      for (int i = 0; i < 3; i++) {
-        if ((extent[i] & 1) == 0)
-          throw Exception("-extent must be a (list of) odd numbers");
-        if (extent[i] > H.size(i))
+      for (ssize_t axis = 0; axis < 3; ++axis) {
+        if (extent[axis] > H.size(axis))
           throw Exception("-extent must not exceed the image dimensions");
+        if ((subsample_factors[axis] & 1) != (extent[axis] & 1))
+          throw Exception("-extent must match subsampling factors "
+                          "(odd for no subsampling or subsampling by an odd factor; "
+                          "even for subsampling by an even factor)");
       }
     } else {
-      uint32_t e = 1;
-      while (Math::pow3(e) < H.size(3))
-        e += 2;
-      extent = {std::min(e, uint32_t(H.size(0))),  //
-                std::min(e, uint32_t(H.size(1))),  //
-                std::min(e, uint32_t(H.size(2)))}; //
+      extent = {subsample_factors[0] & 1 ? 3 : 2, subsample_factors[1] & 1 ? 3 : 2, subsample_factors[2] & 1 ? 3 : 2};
+      ssize_t prev_num_voxels = 0; // Exit loop below if maximum achievable extent is reached
+      while (extent[0] * extent[1] * extent[2] < std::max(H.size(3), prev_num_voxels)) {
+        prev_num_voxels = extent[0] * extent[1] * extent[2];
+        // If multiple axes are tied for spatial extent in mm, increment all of them
+        const default_type min_length =
+            std::min({extent[0] * H.spacing(0), extent[1] * H.spacing(1), extent[2] * H.spacing(2)});
+        for (ssize_t axis = 0; axis != 3; ++axis) {
+          if (extent[axis] * H.spacing(axis) == min_length && extent[axis] + 2 <= H.size(axis))
+            extent[axis] += 2;
+        }
+      }
     }
-    INFO("selected patch size: " + str(extent[0]) + " x " + str(extent[1]) + " x " + str(extent[2]) + ".");
+    INFO("selected cuboid patch size: " + str(extent[0]) + " x " + str(extent[1]) + " x " + str(extent[2]));
 
-    if (std::min<uint32_t>(H.size(3), extent[0] * extent[1] * extent[2]) < 15) {
-      WARN("The number of volumes or the patch size is small. "
-           "This may lead to discretisation effects in the noise level "
-           "and cause inconsistent denoising between adjacent voxels.");
+    if (std::min(H.size(3), extent[0] * extent[1] * extent[2]) < 15) {
+      WARN("The number of volumes or the patch size is small; "
+           "this may lead to discretisation effects in the noise level "
+           "and cause inconsistent denoising between adjacent voxels");
     }
 
-    return std::make_shared<Cuboid>(H, extent);
+    return std::make_shared<Cuboid>(H, extent, subsample_factors);
   } break;
   default:
     assert(false);
diff --git a/src/denoise/kernel/kernel.h b/src/denoise/kernel/kernel.h
index 3bcfb3d2d3..e091074648 100644
--- a/src/denoise/kernel/kernel.h
+++ b/src/denoise/kernel/kernel.h
@@ -16,6 +16,7 @@
 
 #pragma once
 
+#include <array>
 #include <memory>
 #include <string>
 #include <vector>
@@ -29,11 +30,13 @@ namespace MR::Denoise::Kernel {
 class Base;
 
 extern const char *const shape_description;
-extern const char *const size_description;
+extern const char *const default_size_description;
+extern const char *const cuboid_size_description;
 
 const std::vector<std::string> shapes = {"cuboid", "sphere"};
 enum class shape_type { CUBOID, SPHERE };
 extern const App::OptionGroup options;
-std::shared_ptr<Base> make_kernel(const Header &H);
+
+std::shared_ptr<Base> make_kernel(const Header &H, const std::array<ssize_t, 3> &subsample_factors);
 
 } // namespace MR::Denoise::Kernel
diff --git a/src/denoise/kernel/sphere_base.cpp b/src/denoise/kernel/sphere_base.cpp
index 0fadb90354..d441f2e975 100644
--- a/src/denoise/kernel/sphere_base.cpp
+++ b/src/denoise/kernel/sphere_base.cpp
@@ -20,26 +20,59 @@
 
 namespace MR::Denoise::Kernel {
 
-SphereBase::Shared::Shared(const Header &voxel_grid, const default_type max_radius) {
+SphereBase::Shared::Shared(const Header &voxel_grid,
+                           const default_type max_radius,
+                           const std::array<ssize_t, 3> &subsample_factors) {
   const default_type max_radius_sq = Math::pow2(max_radius);
-  const Voxel::index_type half_extents({ssize_t(std::ceil(max_radius / voxel_grid.spacing(0))),   //
-                                        ssize_t(std::ceil(max_radius / voxel_grid.spacing(1))),   //
-                                        ssize_t(std::ceil(max_radius / voxel_grid.spacing(2)))}); //
+  Eigen::Array<int, 3, 2> bounding_box;
+  std::array<default_type, 3> halfvoxel_offsets;
+  for (ssize_t axis = 0; axis != 3; ++axis) {
+    if (subsample_factors[axis] % 2) {
+      bounding_box(axis, 1) = int(std::ceil(max_radius / voxel_grid.spacing(axis)));
+      bounding_box(axis, 0) = -bounding_box(axis, 1);
+      halfvoxel_offsets[axis] = 0.0;
+    } else {
+      bounding_box(axis, 0) = -int(std::ceil((max_radius / voxel_grid.spacing(axis)) - 0.5));
+      bounding_box(axis, 1) = int(std::ceil((max_radius / voxel_grid.spacing(axis)) + 0.5));
+      halfvoxel_offsets[axis] = 0.5;
+    }
+  }
   // Build the searchlight
-  data.reserve(size_t(2 * half_extents[0] + 1) * size_t(2 * half_extents[1] + 1) * size_t(2 * half_extents[2] + 1));
+  data.reserve(size_t(bounding_box(0, 1) + 1 - bounding_box(0, 0)) * //
+               size_t(bounding_box(1, 1) + 1 - bounding_box(1, 0)) * //
+               size_t(bounding_box(2, 1) + 1 - bounding_box(2, 0))); //
   Offset::index_type offset({0, 0, 0});
-  for (offset[2] = -half_extents[2]; offset[2] <= half_extents[2]; ++offset[2]) {
-    for (offset[1] = -half_extents[1]; offset[1] <= half_extents[1]; ++offset[1]) {
-      for (offset[0] = -half_extents[0]; offset[0] <= half_extents[0]; ++offset[0]) {
-        const default_type squared_distance = Math::pow2(offset[0] * voxel_grid.spacing(0))    //
-                                              + Math::pow2(offset[1] * voxel_grid.spacing(1))  //
-                                              + Math::pow2(offset[2] * voxel_grid.spacing(2)); //
+  for (offset[2] = bounding_box(2, 0); offset[2] <= bounding_box(2, 1); ++offset[2]) {
+    for (offset[1] = bounding_box(1, 0); offset[1] <= bounding_box(1, 1); ++offset[1]) {
+      for (offset[0] = bounding_box(0, 0); offset[0] <= bounding_box(0, 1); ++offset[0]) {
+        const default_type squared_distance =
+            Math::pow2((offset[0] + halfvoxel_offsets[0]) * voxel_grid.spacing(0))    //
+            + Math::pow2((offset[1] + halfvoxel_offsets[1]) * voxel_grid.spacing(1))  //
+            + Math::pow2((offset[2] + halfvoxel_offsets[2]) * voxel_grid.spacing(2)); //
         if (squared_distance <= max_radius_sq)
           data.emplace_back(Offset(offset, squared_distance));
       }
     }
   }
   std::sort(data.begin(), data.end());
+  DEBUG("Spherical searchlight construction:");
+  DEBUG("Voxel spacing: ["                   //
+        + str(voxel_grid.spacing(0)) + ","   //
+        + str(voxel_grid.spacing(1)) + ","   //
+        + str(voxel_grid.spacing(2)) + "]"); //
+  DEBUG("Maximum nominated radius: " + str(max_radius));
+  DEBUG("Halfvoxel offsets: ["              //
+        + str(halfvoxel_offsets[0]) + ","   //
+        + str(halfvoxel_offsets[1]) + ","   //
+        + str(halfvoxel_offsets[2]) + "]"); //
+  DEBUG("Bounding box for search: ["        //
+        "[" +
+        str(bounding_box(0, 0)) + " " + str(bounding_box(0, 1)) + "] " +       //
+        "[" + str(bounding_box(1, 0)) + " " + str(bounding_box(1, 1)) + "] " + //
+        "[" + str(bounding_box(2, 0)) + " " + str(bounding_box(2, 1)) + "]]"); //
+  DEBUG("First element: " + str(data.front().index.transpose()) + " @ " + str(data.front().distance()));
+  DEBUG("Last element: " + str(data.back().index.transpose()) + " @ " + str(data.back().distance()));
+  DEBUG("Number of elements: " + str(data.size()));
 }
 
 } // namespace MR::Denoise::Kernel
diff --git a/src/denoise/kernel/sphere_base.h b/src/denoise/kernel/sphere_base.h
index 7dbe7af942..da1bfa4871 100644
--- a/src/denoise/kernel/sphere_base.h
+++ b/src/denoise/kernel/sphere_base.h
@@ -16,6 +16,7 @@
 
 #pragma once
 
+#include <array>
 #include <memory>
 #include <vector>
 
@@ -29,8 +30,8 @@ namespace MR::Denoise::Kernel {
 class SphereBase : public Base {
 
 public:
-  SphereBase(const Header &voxel_grid, const default_type max_radius)
-      : Base(voxel_grid), shared(new Shared(voxel_grid, max_radius)) {}
+  SphereBase(const Header &voxel_grid, const default_type max_radius, const std::array<ssize_t, 3> &subsample_factors)
+      : Base(voxel_grid), shared(new Shared(voxel_grid, max_radius, subsample_factors)) {}
 
   SphereBase(const SphereBase &) = default;
 
@@ -40,7 +41,7 @@ class SphereBase : public Base {
   class Shared {
   public:
     using TableType = std::vector<Offset>;
-    Shared(const Header &voxel_grid, const default_type max_radius);
+    Shared(const Header &voxel_grid, const default_type max_radius, const std::array<ssize_t, 3> &subsample_factors);
     TableType::const_iterator begin() const { return data.begin(); }
     TableType::const_iterator end() const { return data.end(); }
 
diff --git a/src/denoise/kernel/sphere_radius.h b/src/denoise/kernel/sphere_radius.h
index 002dd8243a..73a533d022 100644
--- a/src/denoise/kernel/sphere_radius.h
+++ b/src/denoise/kernel/sphere_radius.h
@@ -25,8 +25,10 @@ namespace MR::Denoise::Kernel {
 
 class SphereFixedRadius : public SphereBase {
 public:
-  SphereFixedRadius(const Header &voxel_grid, const default_type radius)
-      : SphereBase(voxel_grid, radius),                               //
+  SphereFixedRadius(const Header &voxel_grid,                         //
+                    const default_type radius,                        //
+                    const std::array<ssize_t, 3> &subsample_factors)  //
+      : SphereBase(voxel_grid, radius, subsample_factors),            //
         maximum_size(std::distance(shared->begin(), shared->end())) { //
     INFO("Maximum number of voxels in " + str(radius) + "mm fixed-radius kernel is " + str(maximum_size));
   }
diff --git a/src/denoise/kernel/sphere_ratio.h b/src/denoise/kernel/sphere_ratio.h
index bb75268ed6..b6334bc884 100644
--- a/src/denoise/kernel/sphere_ratio.h
+++ b/src/denoise/kernel/sphere_ratio.h
@@ -27,8 +27,8 @@ constexpr default_type sphere_multiplier_default = 1.0 / 0.85;
 class SphereRatio : public SphereBase {
 
 public:
-  SphereRatio(const Header &voxel_grid, const default_type min_ratio)
-      : SphereBase(voxel_grid, compute_max_radius(voxel_grid, min_ratio)),
+  SphereRatio(const Header &voxel_grid, const default_type min_ratio, const std::array<ssize_t, 3> &subsample_factors)
+      : SphereBase(voxel_grid, compute_max_radius(voxel_grid, min_ratio), subsample_factors),
         min_size(std::ceil(voxel_grid.size(3) * min_ratio)) {}
 
   SphereRatio(const SphereRatio &) = default;
diff --git a/src/denoise/recon.cpp b/src/denoise/recon.cpp
index 19ebdaf19b..de48013a05 100644
--- a/src/denoise/recon.cpp
+++ b/src/denoise/recon.cpp
@@ -23,22 +23,28 @@ namespace MR::Denoise {
 template <typename F>
 Recon<F>::Recon(const Header &header,
                 Image<bool> &mask,
+                std::shared_ptr<Subsample> subsample,
                 std::shared_ptr<Kernel::Base> kernel,
                 std::shared_ptr<Estimator::Base> estimator,
                 filter_type filter,
                 aggregator_type aggregator,
                 Exports &exports)
-    : Estimate<F>(header, mask, kernel, estimator, exports),
+    : Estimate<F>(header, mask, subsample, kernel, estimator, exports),
       filter(filter),
       aggregator(aggregator),
-      // FWHM = 2 x cube root of voxel spacings
-      gaussian_multiplier(-std::log(2.0) /
-                          Math::pow2(std::cbrt(header.spacing(0) * header.spacing(1) * header.spacing(2)))),
+      // FWHM = 2 x cube root of spacings between kernels
+      gaussian_multiplier(-std::log(2.0) /                                                           //
+                          Math::pow2(std::cbrt(subsample->get_factors()[0] * header.spacing(0)       //
+                                               * subsample->get_factors()[1] * header.spacing(1)     //
+                                               * subsample->get_factors()[2] * header.spacing(2)))), //
       w(std::min(Estimate<F>::m, kernel->estimated_size())),
       Xr(Estimate<F>::m, aggregator == aggregator_type::EXCLUSIVE ? 1 : kernel->estimated_size()) {}
 
 template <typename F> void Recon<F>::operator()(Image<F> &dwi, Image<F> &out) {
 
+  if (!Estimate<F>::subsample->process({dwi.index(0), dwi.index(1), dwi.index(2)}))
+    return;
+
   Estimate<F>::operator()(dwi);
 
   const ssize_t n = Estimate<F>::neighbourhood.voxels.size();
@@ -66,25 +72,33 @@ template <typename F> void Recon<F>::operator()(Image<F> &dwi, Image<F> &out) {
     sum_weights = double(out_rank);
     break;
   case filter_type::FROBENIUS: {
-    const double beta = r / q;
-    const double transition = 1.0 + std::sqrt(beta);
-    double clam = 0.0;
-    for (ssize_t i = 0; i != r; ++i) {
-      const double lam = std::max(Estimate<F>::s[i], 0.0) / q;
-      clam += lam;
-      const double y = clam / (Estimate<F>::threshold.sigma2 * (i + 1));
-      double nu = 0.0;
-      if (y > transition) {
-        nu = std::sqrt(Math::pow2(Math::pow2(y) - beta - 1.0) - (4.0 * beta)) / y;
-        ++out_rank;
+    if (Estimate<F>::threshold.sigma2 == 0.0) {
+      w.head(r).setOnes();
+      out_rank = r;
+      sum_weights = double(r);
+    } else {
+      const double beta = r / q;
+      const double transition = 1.0 + std::sqrt(beta);
+      double clam = 0.0;
+      for (ssize_t i = 0; i != r; ++i) {
+        const double lam = std::max(Estimate<F>::s[i], 0.0) / q;
+        clam += lam;
+        const double y = clam / (Estimate<F>::threshold.sigma2 * (i + 1));
+        double nu = 0.0;
+        if (y > transition) {
+          nu = std::sqrt(Math::pow2(Math::pow2(y) - beta - 1.0) - (4.0 * beta)) / y;
+          ++out_rank;
+        }
+        w[i] = clam > 0.0 ? (nu / y) : 0.0;
+        sum_weights += w[i];
       }
-      w[i] = nu / y;
-      sum_weights += w[i];
     }
   } break;
   default:
     assert(false);
   }
+  assert(w.head(r).allFinite());
+  assert(std::isfinite(sum_weights));
 
   // recombine data using only eigenvectors above threshold
   // If only the data computed when this voxel was the centre of the patch
@@ -93,7 +107,6 @@ template <typename F> void Recon<F>::operator()(Image<F> &dwi, Image<F> &out) {
   //   if however the result from this patch is to contribute to the synthesized image
   //   for all voxels that were utilised within this patch,
   //   then we need to instead compute the full projection
-  // TODO Use a new data member local to Recon<>
   switch (aggregator) {
   case aggregator_type::EXCLUSIVE:
     if (Estimate<F>::m <= n)
@@ -108,6 +121,7 @@ template <typename F> void Recon<F>::operator()(Image<F> &dwi, Image<F> &out) {
           (Estimate<F>::eig.eigenvectors() *                                                          //
            (w.head(r).cast<F>().matrix().asDiagonal() *                                               //
             Estimate<F>::eig.eigenvectors().adjoint().col(Estimate<F>::neighbourhood.centre_index))); //
+    assert(Xr.allFinite());
     assign_pos_of(dwi).to(out);
     out.row(3) = Xr.col(0);
     if (Estimate<F>::exports.sum_aggregation.valid()) {
@@ -120,18 +134,23 @@ template <typename F> void Recon<F>::operator()(Image<F> &dwi, Image<F> &out) {
     }
     break;
   default: {
-    if (Estimate<F>::m <= n)
-      Xr.noalias() =                                                      //
-          Estimate<F>::eig.eigenvectors() *                               //
-          (w.head(r).cast<F>().matrix().asDiagonal() *                    //
-           (Estimate<F>::eig.eigenvectors().adjoint() * Estimate<F>::X)); //
-    else
+    if (in_rank == r) {
+      Xr.leftCols(n).noalias() = Estimate<F>::X.leftCols(n);
+    } else if (Estimate<F>::m <= n) {
+      Xr.leftCols(n).noalias() =                        //
+          Estimate<F>::eig.eigenvectors() *             //
+          (w.head(r).cast<F>().matrix().asDiagonal() *  //
+           (Estimate<F>::eig.eigenvectors().adjoint() * //
+            Estimate<F>::X.leftCols(n)));               //
+    } else {
       Xr.leftCols(n).noalias() =                         //
           Estimate<F>::X.leftCols(n) *                   //
           (Estimate<F>::eig.eigenvectors() *             //
            (w.head(r).cast<F>().matrix().asDiagonal() *  //
             Estimate<F>::eig.eigenvectors().adjoint())); //
-    std::lock_guard<std::mutex> lock(mutex_aggregator);
+    }
+    assert(Xr.leftCols(n).allFinite());
+    std::lock_guard<std::mutex> lock(Estimate<F>::mutex);
     for (size_t voxel_index = 0; voxel_index != Estimate<F>::neighbourhood.voxels.size(); ++voxel_index) {
       assign_pos_of(Estimate<F>::neighbourhood.voxels[voxel_index].index, 0, 3).to(out);
       assign_pos_of(Estimate<F>::neighbourhood.voxels[voxel_index].index).to(Estimate<F>::exports.sum_aggregation);
@@ -163,8 +182,9 @@ template <typename F> void Recon<F>::operator()(Image<F> &dwi, Image<F> &out) {
   } break;
   }
 
+  auto ss_index = Estimate<F>::subsample->in2ss({dwi.index(0), dwi.index(1), dwi.index(2)});
   if (Estimate<F>::exports.sum_optshrink.valid()) {
-    assign_pos_of(dwi, 0, 3).to(Estimate<F>::exports.sum_optshrink);
+    assign_pos_of(ss_index, 0, 3).to(Estimate<F>::exports.sum_optshrink);
     Estimate<F>::exports.sum_optshrink.value() = sum_weights;
   }
 }
diff --git a/src/denoise/recon.h b/src/denoise/recon.h
index 8646d5b56d..fd48c7ed97 100644
--- a/src/denoise/recon.h
+++ b/src/denoise/recon.h
@@ -18,7 +18,6 @@
 
 #include <limits>
 #include <memory>
-#include <mutex>
 
 #include <Eigen/Dense>
 
@@ -36,6 +35,7 @@ template <typename F> class Recon : public Estimate<F> {
 public:
   Recon(const Header &header,
         Image<bool> &mask,
+        std::shared_ptr<Subsample> subsample,
         std::shared_ptr<Kernel::Base> kernel,
         std::shared_ptr<Estimator::Base> estimator,
         filter_type filter,
@@ -53,13 +53,8 @@ template <typename F> class Recon : public Estimate<F> {
   // Reusable memory
   vector_type w;
   typename Estimate<F>::MatrixType Xr;
-
-  // Some data can only be written in a thread-safe manner
-  static std::mutex mutex_aggregator;
 };
 
-template <typename F> std::mutex Recon<F>::mutex_aggregator;
-
 template class Recon<float>;
 template class Recon<cfloat>;
 template class Recon<double>;
diff --git a/src/denoise/subsample.cpp b/src/denoise/subsample.cpp
new file mode 100644
index 0000000000..ccac67c712
--- /dev/null
+++ b/src/denoise/subsample.cpp
@@ -0,0 +1,108 @@
+/* Copyright (c) 2008-2024 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#include "denoise/subsample.h"
+
+namespace MR::Denoise {
+
+const App::Option subsample_option =
+    App::Option("subsample",
+                "reduce the number of PCA kernels relative to the number of image voxels; "
+                "can provide either an integer subsampling factor, "
+                "or a comma-separated list of three factors;"
+                "default: 2") +
+    App::Argument("factor").type_integer(1);
+
+Subsample::Subsample(const Header &in, const std::array<ssize_t, 3> &factors)
+    : H_in(make_input_header(in)),
+      factors(factors),
+      size({(in.size(0) + factors[0] - 1) / factors[0],
+            (in.size(1) + factors[1] - 1) / factors[1],
+            (in.size(2) + factors[2] - 1) / factors[2]}),
+      origin({(in.size(0) - factors[0] * (size[0] - 1) - 1) / 2,
+              (in.size(1) - factors[1] * (size[1] - 1) - 1) / 2,
+              (in.size(2) - factors[2] * (size[2] - 1) - 1) / 2}),
+      H_ss(make_subsample_header()) {}
+
+bool Subsample::process(const Kernel::Voxel::index_type &pos) const {
+  for (ssize_t axis = 0; axis != 3; ++axis) {
+    if (pos[axis] % factors[axis] != origin[axis])
+      return false;
+  }
+  return true;
+}
+
+std::array<ssize_t, 3> Subsample::in2ss(const Kernel::Voxel::index_type &pos) const {
+  // Do not attempt to map an unprocessed voxel to a voxel index in subsampled space
+  assert(process(pos));
+  assert(!is_out_of_bounds(H_in, pos, 0, 3));
+  return std::array<ssize_t, 3>({(pos[0] - origin[0]) / factors[0],   //
+                                 (pos[1] - origin[1]) / factors[1],   //
+                                 (pos[2] - origin[2]) / factors[2]}); //
+}
+
+std::array<ssize_t, 3> Subsample::ss2in(const Kernel::Voxel::index_type &pos) const {
+  assert(!is_out_of_bounds(H_ss, pos));
+  return std::array<ssize_t, 3>({pos[0] * factors[0] + origin[0],   //
+                                 pos[1] * factors[1] + origin[1],   //
+                                 pos[2] * factors[2] + origin[2]}); //
+}
+
+std::shared_ptr<Subsample> Subsample::make(const Header &in) {
+  auto opt = App::get_options("subsample");
+  std::array<ssize_t, 3> factors({2, 2, 2});
+  if (!opt.empty()) {
+    const std::vector<ssize_t> userinput = parse_ints<ssize_t>(opt[0][0]);
+    if (userinput.size() == 1)
+      factors = {userinput[0], userinput[0], userinput[0]};
+    else if (userinput.size() == 3)
+      factors = {userinput[0], userinput[1], userinput[2]};
+    else
+      throw Exception("Subsampling factor must be either a single positive integer, "
+                      "or a comma-separated list of three positive integers");
+  }
+  return std::make_shared<Subsample>(in, factors);
+}
+
+Header Subsample::make_input_header(const Header &H_in) const {
+  Header H(H_in);
+  H.ndim() = 3;
+  H.reset_intensity_scaling();
+  H.datatype() = DataType::Float32;
+  H.datatype().set_byte_order_native();
+  return H;
+}
+
+Header Subsample::make_subsample_header() const {
+  Header H(H_in);
+  H.ndim() = 3;
+  H.reset_intensity_scaling();
+  H.datatype() = DataType::Float32;
+  H.datatype().set_byte_order_native();
+  for (ssize_t axis = 0; axis != 3; ++axis) {
+    H.size(axis) = size[axis];
+    H.spacing(axis) *= factors[axis];
+  }
+  // Need to move the transform origin from voxel [0,0,0] in the input image
+  //   to voxel [0,0,0] in the subsampled voxel grid;
+  //   this is just a voxel2scanner transformation of "origin"
+  H.transform().translation() = H.transform() * Eigen::Matrix<default_type, 3, 1>({origin[0] * H_in.spacing(0),   //
+                                                                                   origin[1] * H_in.spacing(1),   //
+                                                                                   origin[2] * H_in.spacing(2)}); //
+  return H;
+}
+
+} // namespace MR::Denoise
diff --git a/src/denoise/subsample.h b/src/denoise/subsample.h
new file mode 100644
index 0000000000..f045138748
--- /dev/null
+++ b/src/denoise/subsample.h
@@ -0,0 +1,60 @@
+/* Copyright (c) 2008-2024 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#pragma once
+
+#include <array>
+#include <memory>
+
+#include "app.h"
+#include "denoise/kernel/voxel.h"
+#include "header.h"
+
+namespace MR::Denoise {
+
+extern const App::Option subsample_option;
+
+class Subsample {
+public:
+  Subsample(const Header &in, const std::array<ssize_t, 3> &factors);
+
+  const Header &header() const { return H_ss; }
+
+  // TODO What other functionalities does this class need?
+  // - Decide whether a given 3-vector voxel position should be sampled vs. not
+  // - Convert input image 3-vector voxel position to output subsampled header voxel position
+  // - Convert subsampled header voxel position to centre voxel in input image
+  // TODO May want to move definition of Kernel::Voxel out of Kernel namespace
+  bool process(const Kernel::Voxel::index_type &pos) const;
+  // TODO Rename these
+  std::array<ssize_t, 3> in2ss(const Kernel::Voxel::index_type &pos) const;
+  std::array<ssize_t, 3> ss2in(const Kernel::Voxel::index_type &pos) const;
+  const std::array<ssize_t, 3> &get_factors() const { return factors; }
+
+  static std::shared_ptr<Subsample> make(const Header &in);
+
+protected:
+  const Header H_in;
+  const std::array<ssize_t, 3> factors;
+  const std::array<ssize_t, 3> size;
+  const std::array<ssize_t, 3> origin;
+  const Header H_ss;
+
+  Header make_input_header(const Header &) const;
+  Header make_subsample_header() const;
+};
+
+} // namespace MR::Denoise

From b772559fbbd7d8acc113df069abdab69d1a5f0ba Mon Sep 17 00:00:00 2001
From: Robert Smith <robert.smith@florey.edu.au>
Date: Tue, 26 Nov 2024 15:35:47 +1100
Subject: [PATCH 21/34] dwi*noise: Multiple changes

- Add -estimator med, which provides an estimate of the noise level based on the median eigenvalue as per Gavish and Donoho 2014.
- Perform linear demodulation of complex input data prior to PCA; and in the case of dwidenoise, re-introduce this phase ramp after denoising has completed. A linear phase ramp is determined per k-space group (2D slice / 3D volume for 3D acquisition) based on the maximum of the interpolated k-space data.
- dwidenoise can now filter the eigenspectrum based on noise level truncation (3.0.x behaviour), optimal shrinkage (as implemented previously), or additionally now based on optimal truncation as per Gavish and Donoho 2014.
- mrfilter: New mode "demodulate", which provides a direct interface to the same phase demodulation functionality as is used in dwi*noise.
---
 cmd/dwi2noise.cpp                      |  48 +++++-
 cmd/dwidenoise.cpp                     |  74 ++++++++-
 cmd/mrfilter.cpp                       |  93 +++++++----
 core/filter/demodulate.h               | 222 +++++++++++++++++++++++++
 docs/reference/commands/dwi2noise.rst  |  22 ++-
 docs/reference/commands/dwidenoise.rst |  26 ++-
 docs/reference/commands/mrfilter.rst   |   9 +-
 src/denoise/demodulate.cpp             |  88 ++++++++++
 src/denoise/demodulate.h               |  32 ++++
 src/denoise/denoise.cpp                |   2 +
 src/denoise/denoise.h                  |   6 +-
 src/denoise/estimate.cpp               |   3 +-
 src/denoise/estimator/estimator.cpp    |   5 +
 src/denoise/estimator/estimator.h      |   4 +-
 src/denoise/estimator/med.h            |  56 +++++++
 src/denoise/estimator/result.h         |   2 +-
 src/denoise/recon.cpp                  |  35 +++-
 src/denoise/subsample.cpp              |   2 +-
 18 files changed, 663 insertions(+), 66 deletions(-)
 create mode 100644 core/filter/demodulate.h
 create mode 100644 src/denoise/demodulate.cpp
 create mode 100644 src/denoise/demodulate.h
 create mode 100644 src/denoise/estimator/med.h

diff --git a/cmd/dwi2noise.cpp b/cmd/dwi2noise.cpp
index eda4d4fde6..e22284bc06 100644
--- a/cmd/dwi2noise.cpp
+++ b/cmd/dwi2noise.cpp
@@ -17,13 +17,16 @@
 #include <memory>
 
 #include "algo/threaded_loop.h"
+#include "axes.h"
 #include "command.h"
+#include "denoise/demodulate.h"
 #include "denoise/estimate.h"
 #include "denoise/estimator/estimator.h"
 #include "denoise/exports.h"
 #include "denoise/kernel/kernel.h"
 #include "denoise/subsample.h"
 #include "exception.h"
+#include "filter/demodulate.h"
 
 using namespace MR;
 using namespace App;
@@ -59,6 +62,8 @@ void usage() {
     " the output will be the total noise level across real and imaginary channels,"
     " so a scale factor sqrt(2) applies."
 
+  + demodulation_description
+
   + Kernel::shape_description
 
   + Kernel::default_size_description
@@ -82,7 +87,12 @@ void usage() {
   + "* If using -estimator mrm2022: "
     "Olesen, J.L.; Ianus, A.; Ostergaard, L.; Shemesh, N.; Jespersen, S.N. "
     "Tensor denoising of multidimensional MRI data. "
-    "Magnetic Resonance in Medicine, 2022, 89(3), 1160-1172";
+    "Magnetic Resonance in Medicine, 2022, 89(3), 1160-1172"
+
+  + "* If using -estimator med: "
+    "Gavish, M.; Donoho, D.L."
+    "The Optimal Hard Threshold for Singular Values is 4/sqrt(3). "
+    "IEEE Transactions on Information Theory, 2014, 60(8), 5040-5053.";
 
   ARGUMENTS
   + Argument("dwi", "the input diffusion-weighted image").type_image_in()
@@ -94,6 +104,8 @@ void usage() {
   + Estimator::option
   + Kernel::options
   + subsample_option
+  + demodulation_options
+
 
   // TODO Implement mask option
   // Note that behaviour of -mask for dwi2noise may be different to that of dwidenoise
@@ -154,11 +166,33 @@ void run(Header &data,
   ThreadedLoop("running MP-PCA noise level estimation", data, 0, 3).run(func, input);
 }
 
+template <typename T>
+void run(Header &data,
+         const std::vector<size_t> &demodulation_axes,
+         std::shared_ptr<Subsample> subsample,
+         std::shared_ptr<Kernel::Base> kernel,
+         std::shared_ptr<Estimator::Base> estimator,
+         Exports &exports) {
+  if (demodulation_axes.empty()) {
+    run<T>(data, subsample, kernel, estimator, exports);
+    return;
+  }
+  auto input = data.get_image<T>();
+  auto input_demod = Image<T>::scratch(data, "Phase-demodulated version of \"" + data.name() + "\"");
+  {
+    Filter::Demodulate demodulator(input, demodulation_axes);
+    demodulator(input, input_demod);
+  }
+  Image<bool> mask; // unused
+  Estimate<T> func(data, mask, subsample, kernel, estimator, exports);
+  ThreadedLoop("running MP-PCA noise level estimation", data, 0, 3).run(func, input_demod);
+}
+
 void run() {
   auto dwi = Header::open(argument[0]);
-
   if (dwi.ndim() != 4 || dwi.size(3) <= 1)
     throw Exception("input image must be 4-dimensional");
+  bool complex = dwi.datatype().is_complex();
 
   auto subsample = Subsample::make(dwi);
   assert(subsample);
@@ -184,25 +218,29 @@ void run() {
   if (!opt.empty())
     exports.set_patchcount(opt[0][0]);
 
+  const std::vector<size_t> demodulation_axes = get_demodulation_axes(dwi);
+
   int prec = get_option_value("datatype", 0); // default: single precision
-  if (dwi.datatype().is_complex())
+  if (complex)
     prec += 2; // support complex input data
   switch (prec) {
   case 0:
+    assert(demodulation_axes.empty());
     INFO("select real float32 for processing");
     run<float>(dwi, subsample, kernel, estimator, exports);
     break;
   case 1:
+    assert(demodulation_axes.empty());
     INFO("select real float64 for processing");
     run<double>(dwi, subsample, kernel, estimator, exports);
     break;
   case 2:
     INFO("select complex float32 for processing");
-    run<cfloat>(dwi, subsample, kernel, estimator, exports);
+    run<cfloat>(dwi, demodulation_axes, subsample, kernel, estimator, exports);
     break;
   case 3:
     INFO("select complex float64 for processing");
-    run<cdouble>(dwi, subsample, kernel, estimator, exports);
+    run<cdouble>(dwi, demodulation_axes, subsample, kernel, estimator, exports);
     break;
   }
 }
diff --git a/cmd/dwidenoise.cpp b/cmd/dwidenoise.cpp
index 3a38e9591c..cffcb62dd1 100644
--- a/cmd/dwidenoise.cpp
+++ b/cmd/dwidenoise.cpp
@@ -17,12 +17,15 @@
 #include <string>
 
 #include "command.h"
+#include "filter/demodulate.h"
 #include "header.h"
 #include "image.h"
+#include "stride.h"
 
 #include <Eigen/Dense>
 #include <Eigen/Eigenvalues>
 
+#include "denoise/demodulate.h"
 #include "denoise/denoise.h"
 #include "denoise/estimator/base.h"
 #include "denoise/estimator/estimator.h"
@@ -65,6 +68,8 @@ void usage() {
     " If available, including the MRI phase data can reduce such non-Gaussian biases,"
     " and the command now supports complex input data."
 
+  + demodulation_description
+
   + Kernel::shape_description
 
   + Kernel::default_size_description
@@ -76,6 +81,8 @@ void usage() {
     "Hard truncation of sub-threshold components and inclusion of supra-threshold components"
     "---which was the behaviour of the dwidenoise command in version 3.0.x---"
     "can be activated using -filter truncate."
+    "Alternatively, optimal truncation as described in Gavish and Donoho 2014 "
+    "can be utilised by specifying -filter optthresh."
 
   + "-aggregation exclusive corresponds to the behaviour of the dwidenoise command in version 3.0.x, "
     "where the output intensities for a given image voxel are determined exclusively "
@@ -121,7 +128,12 @@ void usage() {
   + "* If using anything other than -aggregation exclusive: "
     "Manjon, J.V.; Coupe, P.; Concha, L.; Buades, A.; D. Collins, D.L.; Robles, M. "
     "Diffusion Weighted Image Denoising Using Overcomplete Local PCA. "
-    "PLoS ONE, 2013, 8(9), e73021";
+    "PLoS ONE, 2013, 8(9), e73021"
+
+  + "* If using -estimator med or -filter optthresh: "
+    "Gavish, M.; Donoho, D.L."
+    "The Optimal Hard Threshold for Singular Values is 4/sqrt(3). "
+    "IEEE Transactions on Information Theory, 2014, 60(8), 5040-5053.";
 
   ARGUMENTS
   + Argument("dwi", "the input diffusion-weighted image.").type_image_in()
@@ -133,6 +145,7 @@ void usage() {
   + Estimator::option
   + Kernel::options
   + subsample_option
+  + demodulation_options
 
   + OptionGroup("Options that affect reconstruction of the output image series")
   // TODO Separate masks for voxels to contribute to patches vs. voxels for which to perform denoising
@@ -143,7 +156,7 @@ void usage() {
            "Modulate how component contributions are filtered "
            "based on the cumulative eigenvalues relative to the noise level; "
            "options are: " + join(filters, ",") + "; "
-           "default: frobenius (Optimal Shrinkage based on minimisation of the Frobenius norm)")
+           "default: optshrink (Optimal Shrinkage based on minimisation of the Frobenius norm)")
     + Argument("choice").type_choice(filters)
   + Option("aggregator",
            "Select how the outcomes of multiple PCA outcomes centred at different voxels "
@@ -270,6 +283,53 @@ void run(Header &data,
   }
 }
 
+template <typename T>
+void run(Header &data,
+         Image<bool> &mask,
+         const std::vector<size_t> &demodulation_axes,
+         std::shared_ptr<Subsample> subsample,
+         std::shared_ptr<Kernel::Base> kernel,
+         std::shared_ptr<Estimator::Base> estimator,
+         filter_type filter,
+         aggregator_type aggregator,
+         const std::string &output_name,
+         Exports &exports) {
+  if (demodulation_axes.empty()) {
+    run<T>(data, mask, subsample, kernel, estimator, filter, aggregator, output_name, exports);
+    return;
+  }
+  auto input = data.get_image<T>();
+  // generate scratch version of DWI with phase demodulation
+  Header H_scratch(data);
+  Stride::set(H_scratch, Stride::contiguous_along_axis(3));
+  H_scratch.datatype() = DataType::from<T>();
+  H_scratch.datatype().set_byte_order_native();
+  auto input_demodulated = Image<T>::scratch(H_scratch, "Phase-demodulated version of input DWI");
+  Filter::Demodulate demodulate(input, demodulation_axes);
+  demodulate(input, input_demodulated, false);
+  input = Image<T>(); // free memory
+  // create output
+  Header header(data);
+  header.datatype() = DataType::from<T>();
+  auto output = Image<T>::create(output_name, header);
+  // run
+  Recon<T> func(data, mask, subsample, kernel, estimator, filter, aggregator, exports);
+  ThreadedLoop("running MP-PCA denoising", data, 0, 3).run(func, input_demodulated, output);
+  // Re-apply phase ramps that were previously demodulated
+  demodulate(output, true);
+  // Rescale output if performing aggregation
+  if (aggregator == aggregator_type::EXCLUSIVE)
+    return;
+  for (auto l_voxel = Loop(exports.sum_aggregation)(output, exports.sum_aggregation); l_voxel; ++l_voxel) {
+    for (auto l_volume = Loop(3)(output); l_volume; ++l_volume)
+      output.value() /= float(exports.sum_aggregation.value());
+  }
+  if (exports.rank_output.valid()) {
+    for (auto l = Loop(exports.sum_aggregation)(exports.rank_output, exports.sum_aggregation); l; ++l)
+      exports.rank_output.value() /= exports.sum_aggregation.value();
+  }
+}
+
 void run() {
   auto dwi = Header::open(argument[0]);
 
@@ -292,7 +352,7 @@ void run() {
   auto estimator = Estimator::make_estimator();
   assert(estimator);
 
-  filter_type filter = filter_type::FROBENIUS;
+  filter_type filter = filter_type::OPTSHRINK;
   opt = get_options("filter");
   if (!opt.empty())
     filter = filter_type(int(opt[0][0]));
@@ -352,25 +412,29 @@ void run() {
     exports.set_sum_aggregation("");
   }
 
+  const std::vector<size_t> demodulation_axes = get_demodulation_axes(dwi);
+
   int prec = get_option_value("datatype", 0); // default: single precision
   if (dwi.datatype().is_complex())
     prec += 2; // support complex input data
   switch (prec) {
   case 0:
+    assert(demodulation_axes.empty());
     INFO("select real float32 for processing");
     run<float>(dwi, mask, subsample, kernel, estimator, filter, aggregator, argument[1], exports);
     break;
   case 1:
+    assert(demodulation_axes.empty());
     INFO("select real float64 for processing");
     run<double>(dwi, mask, subsample, kernel, estimator, filter, aggregator, argument[1], exports);
     break;
   case 2:
     INFO("select complex float32 for processing");
-    run<cfloat>(dwi, mask, subsample, kernel, estimator, filter, aggregator, argument[1], exports);
+    run<cfloat>(dwi, mask, demodulation_axes, subsample, kernel, estimator, filter, aggregator, argument[1], exports);
     break;
   case 3:
     INFO("select complex float64 for processing");
-    run<cdouble>(dwi, mask, subsample, kernel, estimator, filter, aggregator, argument[1], exports);
+    run<cdouble>(dwi, mask, demodulation_axes, subsample, kernel, estimator, filter, aggregator, argument[1], exports);
     break;
   }
 }
diff --git a/cmd/mrfilter.cpp b/cmd/mrfilter.cpp
index e97e2882a4..4e37def026 100644
--- a/cmd/mrfilter.cpp
+++ b/cmd/mrfilter.cpp
@@ -15,28 +15,33 @@
  */
 
 #include <complex>
+#include <set>
 
 #include "command.h"
 #include "filter/base.h"
+#include "filter/demodulate.h"
 #include "filter/gradient.h"
 #include "filter/median.h"
 #include "filter/normalise.h"
 #include "filter/smooth.h"
 #include "filter/zclean.h"
 #include "image.h"
+#include "interp/cubic.h"
 #include "math/fft.h"
 
 using namespace MR;
 using namespace App;
 
-const std::vector<std::string> filters = {"fft", "gradient", "median", "smooth", "normalise", "zclean"};
+const std::vector<std::string> filters = {"demodulate", "fft", "gradient", "median", "smooth", "normalise", "zclean"};
 
 // clang-format off
-const OptionGroup FFTOption = OptionGroup ("Options for FFT filter")
+const OptionGroup FFTAxesOption = OptionGroup ("Options applicable to both demodulate and FFT filters")
   + Option ("axes", "the axes along which to apply the Fourier Transform."
                     " By default, the transform is applied along the three spatial axes."
                     " Provide as a comma-separate list of axis indices.")
-    + Argument ("list").type_sequence_int()
+    + Argument ("list").type_sequence_int();
+
+const OptionGroup FFTOption = OptionGroup ("Options for FFT filter")
   + Option ("inverse", "apply the inverse FFT")
   + Option ("magnitude", "output a magnitude image rather than a complex-valued image")
   + Option ("rescale", "rescale values so that inverse FFT recovers original values")
@@ -116,7 +121,7 @@ void usage() {
 
   DESCRIPTION
   + "The available filters are:"
-    " fft, gradient, median, smooth, normalise, zclean."
+    " demodulate, fft, gradient, median, smooth, normalise, zclean."
   + "Each filter has its own unique set of optional parameters."
   + "For 4D images, each 3D volume is processed independently.";
 
@@ -126,6 +131,7 @@ void usage() {
   + Argument ("output", "the output image.").type_image_out ();
 
   OPTIONS
+  + FFTAxesOption
   + FFTOption
   + GradientOption
   + MedianOption
@@ -137,44 +143,69 @@ void usage() {
 }
 // clang-format on
 
+std::vector<size_t> get_axes(const Header &H, const std::vector<size_t> &default_axes) {
+  auto opt = get_options("axes");
+  std::vector<size_t> axes = default_axes;
+  if (!opt.empty()) {
+    axes = parse_ints<size_t>(opt[0][0]);
+    for (const auto axis : axes)
+      if (axis >= H.ndim())
+        throw Exception("axis provided with -axes option is out of range");
+    if (std::set<size_t>(axes.begin(), axes.end()).size() != axes.size())
+      throw Exception("axis indices must not contain duplicates");
+  }
+  return axes;
+}
+
 void run() {
 
   const size_t filter_index = argument[1];
 
   switch (filter_index) {
 
-  // FFT
+  // Phase demodulation
   case 0: {
-    // FIXME Had to use cdouble throughout; seems to fail at compile time even trying to
-    //   convert between cfloat and cdouble...
-    auto input = Image<cdouble>::open(argument[0]);
+    Header H_in = Header::open(argument[0]);
+    if (!H_in.datatype().is_complex())
+      throw Exception("demodulation filter only applicable for complex image data");
+    auto input = H_in.get_image<cdouble>();
 
-    std::vector<size_t> axes = {0, 1, 2};
-    auto opt = get_options("axes");
-    if (!opt.empty()) {
-      axes = parse_ints<size_t>(opt[0][0]);
-      for (const auto axis : axes)
-        if (axis >= input.ndim())
-          throw Exception("axis provided with -axes option is out of range");
-    }
+    const std::vector<size_t> inner_axes = get_axes(H_in, {0, 1});
+
+    Filter::Demodulate filter(input, inner_axes);
+
+    Header H_out(H_in);
+    Stride::set_from_command_line(H_out);
+    auto output = Image<cdouble>::create(argument[2], H_out);
+
+    filter(input, output);
+  } break;
+
+  // FFT
+  case 1: {
+    Header H_in = Header::open(argument[0]);
+    std::vector<size_t> axes = get_axes(H_in, {0, 1, 2});
     const int direction = get_options("inverse").empty() ? FFTW_FORWARD : FFTW_BACKWARD;
     const bool centre_FFT = !get_options("centre_zero").empty();
     const bool magnitude = !get_options("magnitude").empty();
 
-    Header header = input;
-    Stride::set_from_command_line(header);
-    header.datatype() = magnitude ? DataType::Float32 : DataType::CFloat64;
-    auto output = Image<cdouble>::create(argument[2], header);
+    Header H_out(H_in);
+    Stride::set_from_command_line(H_out);
+    H_out.datatype() = magnitude ? DataType::Float32 : DataType::CFloat64;
+    auto output = Image<cdouble>::create(argument[2], H_out);
     double scale = 1.0;
 
+    // FIXME Had to use cdouble throughout; seems to fail at compile time even trying to
+    //   convert between cfloat and cdouble...
+    auto input = H_in.get_image<cdouble>();
+
     Image<cdouble> in(input), out;
     for (size_t n = 0; n < axes.size(); ++n) {
       scale *= in.size(axes[n]);
-      if (n >= (axes.size() - 1) && !magnitude) {
+      if (n == (axes.size() - 1) && !magnitude) {
         out = output;
-      } else {
-        if (!out.valid())
-          out = Image<cdouble>::scratch(input);
+      } else if (!out.valid()) {
+        out = Image<cdouble>::scratch(H_in);
       }
 
       Math::FFT(in, out, axes[n], direction, centre_FFT);
@@ -187,15 +218,15 @@ void run() {
           [](decltype(out) &a, decltype(output) &b) { a.value() = abs(cdouble(b.value())); }, output, out);
     }
     if (!get_options("rescale").empty()) {
-      scale = std::sqrt(scale);
-      ThreadedLoop(out).run([&scale](decltype(out) &a) { a.value() /= scale; }, output);
+      scale = 1.0 / std::sqrt(scale);
+      ThreadedLoop(out).run([&scale](decltype(out) &a) { a.value() *= scale; }, output);
     }
 
     break;
   }
 
   // Gradient
-  case 1: {
+  case 2: {
     auto input = Image<float>::open(argument[0]);
     Filter::Gradient filter(input, !get_options("magnitude").empty());
 
@@ -224,7 +255,7 @@ void run() {
   }
 
   // Median
-  case 2: {
+  case 3: {
     auto input = Image<float>::open(argument[0]);
     Filter::Median filter(input);
 
@@ -241,7 +272,7 @@ void run() {
   }
 
   // Smooth
-  case 3: {
+  case 4: {
     auto input = Image<float>::open(argument[0]);
     Filter::Smooth filter(input);
 
@@ -272,7 +303,7 @@ void run() {
   }
 
   // Normalisation
-  case 4: {
+  case 5: {
     auto input = Image<float>::open(argument[0]);
     Filter::Normalise filter(input);
 
@@ -289,7 +320,7 @@ void run() {
   }
 
   // Zclean
-  case 5: {
+  case 6: {
     auto input = Image<float>::open(argument[0]);
     Filter::ZClean filter(input);
 
diff --git a/core/filter/demodulate.h b/core/filter/demodulate.h
new file mode 100644
index 0000000000..d28aad4741
--- /dev/null
+++ b/core/filter/demodulate.h
@@ -0,0 +1,222 @@
+/* Copyright (c) 2008-2024 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#pragma once
+
+#include <vector>
+
+#include "algo/copy.h"
+#include "algo/loop.h"
+#include "filter/base.h"
+#include "filter/smooth.h"
+#include "image.h"
+#include "interp/cubic.h"
+#include "math/fft.h"
+#include "progressbar.h"
+
+namespace MR::Filter {
+/** \addtogroup Filters
+@{ */
+
+/*! Estimate a linear phase ramp of a complex image and demodulate by such
+ */
+class Demodulate : public Base {
+public:
+  template <class ImageType>
+  Demodulate(ImageType &in, const std::vector<size_t> &inner_axes)
+      : Base(in), phase(Image<cfloat>::scratch(in, "Scratch image storing linear phase for demodulator")) {
+
+    using value_type = typename ImageType::value_type;
+    using real_type = typename ImageType::value_type::value_type;
+
+    ImageType input(in);
+    // if (!in.datatype().is_complex())
+    //   throw Exception("demodulation filter only applicable for complex image data");
+
+    std::vector<size_t> outer_axes;
+    std::vector<size_t>::const_iterator it = inner_axes.begin();
+    for (size_t axis = 0; axis != input.ndim(); ++axis) {
+      if (it != inner_axes.end() && *it == axis)
+        ++it;
+      else
+        outer_axes.push_back(axis);
+    }
+
+    ProgressBar progress("estimating linear phase modulation", inner_axes.size() + 1);
+
+    // FFT currently hard-wired to double precision;
+    //   have to manually load into cdouble memory
+    Image<cdouble> kspace;
+    {
+      Image<cdouble> temp;
+      for (ssize_t n = 0; n != inner_axes.size(); ++n) {
+        switch (n) {
+        case 0:
+          // TODO For now doing a straight copy;
+          //   would be preferable if, if the type of the input image is cdouble,
+          //   to just make temp a view of input
+          // In retrospect it's probably necessary to template the whole k-space derivation
+          temp = Image<cdouble>::scratch(in, "Scratch k-space for phase demodulator (1)");
+          copy(input, temp);
+          kspace = Image<cdouble>::scratch(in, "Scratch k-space for phase demodulator (2)");
+          break;
+        // TODO If it's possible to prevent the Algo::copy() call above,
+        //   then the below code block can be re-introduced
+        // case 1:
+        // temp = kspace;
+        // kspace = Image<cdouble>::scratch(in, "Scratch k-space for phase demodulator (2)");
+        // break;
+        default:
+          std::swap(temp, kspace);
+          break;
+        }
+        Math::FFT(temp, kspace, inner_axes[n], FFTW_FORWARD, true);
+        ++progress;
+      }
+    }
+
+    auto gen_phase =
+        [&](Image<value_type> &input, Image<cdouble> &kspace, Image<cfloat> &phase, const std::vector<size_t> &axes) {
+          std::array<bool, 3> axis_mask({false, false, false});
+          for (auto axis : axes) {
+            if (axis > 2)
+              throw Exception("Linear phase demodulator does not support non-spatial inner axes");
+            axis_mask[axis] = true;
+          }
+
+          std::vector<ssize_t> index_of_max(kspace.ndim());
+          std::vector<size_t>::const_iterator it = axes.begin();
+          for (ssize_t axis = 0; axis != kspace.ndim(); ++axis) {
+            if (it != axes.end() && *it == axis) {
+              index_of_max[axis] = -1;
+              ++it;
+            } else {
+              index_of_max[axis] = kspace.index(axis);
+            }
+          }
+          double max_magnitude = -std::numeric_limits<double>::infinity();
+          cdouble value_at_max = cdouble(std::numeric_limits<double>::signaling_NaN(),  //
+                                         std::numeric_limits<double>::signaling_NaN()); //
+          for (auto l_inner = Loop(axes)(kspace); l_inner; ++l_inner) {
+            const cdouble value = kspace.value();
+            const double magnitude = std::abs(value);
+            if (magnitude > max_magnitude) {
+              max_magnitude = magnitude;
+              value_at_max = value;
+              for (auto axis : axes)
+                index_of_max[axis] = kspace.index(axis);
+            }
+          }
+
+          using pos_type = Eigen::Matrix<real_type, 3, 1>;
+          using gradient_type = Eigen::Matrix<cdouble, 1, 3>;
+          pos_type pos({real_type(index_of_max[0]), real_type(index_of_max[1]), real_type(index_of_max[2])});
+          gradient_type gradient;
+          cdouble value;
+          Interp::SplineInterp<Image<cdouble>,
+                               Math::UniformBSpline<cdouble>,
+                               Math::SplineProcessingType::ValueAndDerivative>
+              interp(kspace);
+          assign_pos_of(index_of_max).to(interp);
+          interp.voxel(pos.template cast<double>());
+          interp.value_and_gradient(value, gradient);
+          const double mag = std::abs(value);
+          for (ssize_t iter = 0; iter != 10; ++iter) {
+            pos_type grad_mag({0.0, 0.0, 0.0}); // Gradient of the magnitude of the complex k-space data
+            for (ssize_t axis = 0; axis != 3; ++axis) {
+              if (axis_mask[axis])
+                grad_mag[axis] = (value.real() * gradient[axis].real() + value.imag() * gradient[axis].imag()) / mag;
+            }
+            grad_mag /= max_magnitude;
+            pos += 1.0 * grad_mag;
+            interp.value_and_gradient(value, gradient);
+          }
+          value_at_max = value;
+          const real_type phase_at_max = std::arg(value_at_max);
+          max_magnitude = std::abs(value_at_max);
+
+          // Determine direction and frequency of harmonic
+          pos_type kspace_origin;
+          for (ssize_t axis = 0; axis != 3; ++axis)
+            // Do integer division, then convert to floating-point
+            kspace_origin[axis] = axis_mask[axis] ? real_type((kspace.size(axis) - 1) / 2) : real_type(0);
+          const pos_type kspace_offset({axis_mask[0] ? (pos[0] - kspace_origin[0]) : real_type(0),
+                                        axis_mask[1] ? (pos[1] - kspace_origin[1]) : real_type(0),
+                                        axis_mask[2] ? (pos[2] - kspace_origin[2]) : real_type(0)});
+          const pos_type cycle_voxel({axis_mask[0] ? (input.size(0) / kspace_offset[0]) : real_type(0),
+                                      axis_mask[1] ? (input.size(1) / kspace_offset[1]) : real_type(0),
+                                      axis_mask[2] ? (input.size(2) / kspace_offset[2]) : real_type(0)});
+          const pos_type voxel_origin({real_type(0.5 * (input.size(0) - 1.0)),
+                                       real_type(0.5 * (input.size(1) - 1.0)),
+                                       real_type(0.5 * (input.size(2) - 1.0))});
+          for (auto l = Loop(axes)(input, phase); l; ++l) {
+            pos = {real_type(input.index(0)), real_type(input.index(1)), real_type(input.index(2))};
+            const real_type voxel_phase =
+                phase_at_max + (2.0 * Math::pi * (pos - voxel_origin).dot(cycle_voxel) / cycle_voxel.squaredNorm());
+            const value_type modulator(std::cos(voxel_phase), std::sin(voxel_phase));
+            phase.value() = cfloat(modulator);
+          }
+        };
+
+    if (outer_axes.size()) {
+      // TODO Multi-thread
+      // ThreadedLoop("Estimating phase ramps", input, outer_axes).run(gen_phase, input, kspace);
+      for (auto l_outer = Loop(outer_axes)(input, kspace, phase); l_outer; ++l_outer)
+        gen_phase(input, kspace, phase, inner_axes);
+    } else {
+      gen_phase(input, kspace, phase, inner_axes);
+    }
+  }
+
+  template <class InputImageType, class OutputImageType>
+  void operator()(InputImageType &in, OutputImageType &out, const bool forward = false) {
+    // if (!out.datatype().is_complex())
+    //   throw Exception("Cannot modulate phase for output image that is not complex");
+    if (forward) {
+      for (auto l = Loop("re-applying phase modulation")(in, phase, out); l; ++l)
+        out.value() = typename OutputImageType::value_type(typename InputImageType::value_type(in.value()) *
+                                                           typename InputImageType::value_type(cfloat(phase.value())));
+    } else {
+      for (auto l = Loop("performing phase demodulation")(in, phase, out); l; ++l)
+        out.value() =
+            typename OutputImageType::value_type(typename InputImageType::value_type(in.value()) *
+                                                 typename InputImageType::value_type(std::conj(cfloat(phase.value()))));
+    }
+  }
+
+  template <class ImageType> void operator()(ImageType &image, const bool forward = false) {
+    // if (!image.datatype().is_complex())
+    //   throw Exception("Cannot modulate phase for output image that is not complex");
+    if (forward) {
+      for (auto l = Loop("re-applying phase modulation")(image, phase); l; ++l)
+        image.value() = typename ImageType::value_type(typename ImageType::value_type(image.value()) *
+                                                       typename ImageType::value_type(cfloat(phase.value())));
+    } else {
+      for (auto l = Loop("performing phase demodulation")(image, phase); l; ++l)
+        image.value() =
+            typename ImageType::value_type(typename ImageType::value_type(image.value()) *
+                                           typename ImageType::value_type(std::conj(cfloat(phase.value()))));
+    }
+  }
+
+protected:
+  // TODO Change to Image<float>; can produce complex value at processing time
+  Image<cfloat> phase;
+
+  // TODO Define a protected class that can be utilised to generate the phase image in a multi-threaded manner
+};
+//! @}
+} // namespace MR::Filter
diff --git a/docs/reference/commands/dwi2noise.rst b/docs/reference/commands/dwi2noise.rst
index cf4a1f1383..5bee914e21 100644
--- a/docs/reference/commands/dwi2noise.rst
+++ b/docs/reference/commands/dwi2noise.rst
@@ -21,7 +21,7 @@ Usage
 Description
 -----------
 
-DWI data noise map estimation by exploiting data redundancy in the PCA domain using the prior knowledge that the eigenspectrum of random covariance matrices is described by the universal Marchenko-Pastur (MP) distribution. Fitting the MP distribution to the spectrum of patch-wise signal matrices hence provides an estimator of the noise level 'sigma'.
+DWI data noise map estimation by interrogating data redundancy in the PCA domain using the prior knowledge that the eigenspectrum of random covariance matrices is described by the universal Marchenko-Pastur (MP) distribution. Fitting the MP distribution to the spectrum of patch-wise signal matrices hence provides an estimator of the noise level 'sigma'.
 
 Unlike the MRtrix3 command dwidenoise, this command does not generate a denoised version of the input image series; its primary output is instead a map of the estimated noise level. While this can also be obtained from the dwidenoise command using option -noise_out, using instead the dwi2noise command gives the ability to obtain a noise map to which filtering can be applied, which can then be utilised for the actual image series denoising, without generating an unwanted intermiedate denoised image series.
 
@@ -29,10 +29,14 @@ Important note: noise level estimation should only be performed as the first ste
 
 Note that on complex input data, the output will be the total noise level across real and imaginary channels, so a scale factor sqrt(2) applies.
 
+If the input data are of complex type, then a linear phase term will be removed from each k-space prior to PCA. In the absence of metadata indicating otherwise, it is inferred that the first two axes correspond to acquired slices, and different slices / volumes will be demodulated individually; this behaviour can be modified using the -demod_axes option.
+
 The sliding spatial window behaves differently at the edges of the image FoV depending on the shape / size selected for that window. The default behaviour is to use a spherical kernel centred at the voxel of interest, whose size is some multiple of the number of input volumes; where some such voxels lie outside of the image FoV, the radius of the kernel will be increased until the requisite number of voxels are used. For a spherical kernel of a fixed radius, no such expansion will occur, and so for voxels near the image edge a reduced number of voxels will be present in the kernel. For a cuboid kernel, the centre of the kernel will be offset from the voxel being processed such that the entire volume of the kernel resides within the image FoV.
 
 The size of the default spherical kernel is set to select a number of voxels that is 1.0 / 0.85 ~ 1.18 times the number of volumes in the input series. If a cuboid kernel is requested, but the -extent option is not specified, the command will select the smallest isotropic patch size that exceeds the number of DW images in the input data; e.g., 5x5x5 for data with <= 125 DWI volumes, 7x7x7 for data with <= 343 DWI volumes, etc.
 
+Permissible sizes for the cuboid kernel depend on the subsampling factor. If no subsampling is performed, or the subsampling factor is odd, then the extent(s) of the kernel must be odd, such that a unique voxel lies at the very centre of each kernel. If however an even subsampling factor is used, then the extent(s) of the kernel must be even, reflecting the fact that it is a voxel corner that resides at the centre of the kernel.In either case, if the extent is specified manually, the user can either provide a single integer---which will determine the number of voxels in the kernel across all three spatial axes---or a comma-separated list of three integers,individually defining the number of voxels in the kernel for all three spatial axes.
+
 Options
 -------
 
@@ -44,6 +48,7 @@ Options for modifying PCA computations
 -  **-estimator algorithm** Select the noise level estimator (default = Exp2), either:  |br|
    * Exp1: the original estimator used in Veraart et al. (2016);  |br|
    * Exp2: the improved estimator introduced in Cordero-Grande et al. (2019);  |br|
+   * Med: estimate based on the median eigenvalue as in Gavish and Donohue (2014);  |br|
    * MRM2022: the alternative estimator introduced in Olesen et al. (2022).
 
 Options for controlling the sliding spatial window kernel
@@ -55,7 +60,16 @@ Options for controlling the sliding spatial window kernel
 
 -  **-radius_ratio value** Set the spherical kernel size as a ratio of number of voxels to number of input volumes (default: 1.0/0.85 ~= 1.18)
 
--  **-extent window** Set the patch size of the cuboid kernel; can be either a single odd integer or a comma-separated triplet of odd integers
+-  **-extent window** Set the patch size of the cuboid kernel; can be either a single integer or a comma-separated triplet of integers (see Description)
+
+-  **-subsample factor** reduce the number of PCA kernels relative to the number of image voxels; can provide either an integer subsampling factor, or a comma-separated list of three factors;default: 2
+
+Options for phase demodulation of complex data
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+-  **-nodemod** disable phase demodulation
+
+-  **-demod_axes axes** comma-separated list of axis indices along which FFT can be applied for phase demodulation
 
 Options for exporting additional data regarding PCA behaviour
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
@@ -69,6 +83,8 @@ Options for debugging the operation of sliding window kernels
 
 -  **-voxelcount image** The number of voxels that contributed to the PCA for processing of each voxel
 
+-  **-patchcount image** The number of unique patches to which an image voxel contributes
+
 Standard options
 ^^^^^^^^^^^^^^^^
 
@@ -97,6 +113,8 @@ Cordero-Grande, L.; Christiaens, D.; Hutter, J.; Price, A.N.; Hajnal, J.V. Compl
 
 * If using -estimator mrm2022: Olesen, J.L.; Ianus, A.; Ostergaard, L.; Shemesh, N.; Jespersen, S.N. Tensor denoising of multidimensional MRI data. Magnetic Resonance in Medicine, 2022, 89(3), 1160-1172
 
+* If using -estimator med: Gavish, M.; Donoho, D.L.The Optimal Hard Threshold for Singular Values is 4/sqrt(3). IEEE Transactions on Information Theory, 2014, 60(8), 5040-5053.
+
 Tournier, J.-D.; Smith, R. E.; Raffelt, D.; Tabbara, R.; Dhollander, T.; Pietsch, M.; Christiaens, D.; Jeurissen, B.; Yeh, C.-H. & Connelly, A. MRtrix3: A fast, flexible and open software framework for medical image processing and visualisation. NeuroImage, 2019, 202, 116137
 
 --------------
diff --git a/docs/reference/commands/dwidenoise.rst b/docs/reference/commands/dwidenoise.rst
index c769924859..581c8ffc08 100644
--- a/docs/reference/commands/dwidenoise.rst
+++ b/docs/reference/commands/dwidenoise.rst
@@ -27,11 +27,15 @@ Important note: image denoising must be performed as the first step of the image
 
 Note that this function does not correct for non-Gaussian noise biases present in magnitude-reconstructed MRI images. If available, including the MRI phase data can reduce such non-Gaussian biases, and the command now supports complex input data.
 
+If the input data are of complex type, then a linear phase term will be removed from each k-space prior to PCA. In the absence of metadata indicating otherwise, it is inferred that the first two axes correspond to acquired slices, and different slices / volumes will be demodulated individually; this behaviour can be modified using the -demod_axes option.
+
 The sliding spatial window behaves differently at the edges of the image FoV depending on the shape / size selected for that window. The default behaviour is to use a spherical kernel centred at the voxel of interest, whose size is some multiple of the number of input volumes; where some such voxels lie outside of the image FoV, the radius of the kernel will be increased until the requisite number of voxels are used. For a spherical kernel of a fixed radius, no such expansion will occur, and so for voxels near the image edge a reduced number of voxels will be present in the kernel. For a cuboid kernel, the centre of the kernel will be offset from the voxel being processed such that the entire volume of the kernel resides within the image FoV.
 
 The size of the default spherical kernel is set to select a number of voxels that is 1.0 / 0.85 ~ 1.18 times the number of volumes in the input series. If a cuboid kernel is requested, but the -extent option is not specified, the command will select the smallest isotropic patch size that exceeds the number of DW images in the input data; e.g., 5x5x5 for data with <= 125 DWI volumes, 7x7x7 for data with <= 343 DWI volumes, etc.
 
-By default, optimal value shrinkage based on minimisation of the Frobenius norm will be used to attenuate eigenvectors based on the estimated noise level. Hard truncation of sub-threshold components and inclusion of supra-threshold components---which was the behaviour of the dwidenoise command in version 3.0.x---can be activated using -filter truncate.
+Permissible sizes for the cuboid kernel depend on the subsampling factor. If no subsampling is performed, or the subsampling factor is odd, then the extent(s) of the kernel must be odd, such that a unique voxel lies at the very centre of each kernel. If however an even subsampling factor is used, then the extent(s) of the kernel must be even, reflecting the fact that it is a voxel corner that resides at the centre of the kernel.In either case, if the extent is specified manually, the user can either provide a single integer---which will determine the number of voxels in the kernel across all three spatial axes---or a comma-separated list of three integers,individually defining the number of voxels in the kernel for all three spatial axes.
+
+By default, optimal value shrinkage based on minimisation of the Frobenius norm will be used to attenuate eigenvectors based on the estimated noise level. Hard truncation of sub-threshold components and inclusion of supra-threshold components---which was the behaviour of the dwidenoise command in version 3.0.x---can be activated using -filter truncate.Alternatively, optimal truncation as described in Gavish and Donoho 2014 can be utilised by specifying -filter optthresh.
 
 -aggregation exclusive corresponds to the behaviour of the dwidenoise command in version 3.0.x, where the output intensities for a given image voxel are determined exclusively from the PCA decomposition where the sliding spatial window is centred at that voxel. In all other use cases, so-called "overcomplete local PCA" is performed, where the intensities for an output image voxel are some combination of all PCA decompositions for which that voxel is included in the local spatial kernel. There are multiple algebraic forms that modulate the weight with which each decomposition contributes with greater or lesser strength toward the output image intensities. The various options are: 'gaussian': A Gaussian distribution with FWHM equal to twice the voxel size, such that decompisitions centred more closely to the output voxel have greater influence; 'invl0': The inverse of the L0 norm (ie. rank) of each decomposition, as used in Manjon et al. 2013; 'rank': The rank of each decomposition, such that high-rank decompositions contribute more strongly to the output intensities regardless of distance between the output voxel and the centre of the decomposition kernel; 'uniform': All decompositions that include the output voxel in the sliding spatial window contribute equally.
 
@@ -46,6 +50,7 @@ Options for modifying PCA computations
 -  **-estimator algorithm** Select the noise level estimator (default = Exp2), either:  |br|
    * Exp1: the original estimator used in Veraart et al. (2016);  |br|
    * Exp2: the improved estimator introduced in Cordero-Grande et al. (2019);  |br|
+   * Med: estimate based on the median eigenvalue as in Gavish and Donohue (2014);  |br|
    * MRM2022: the alternative estimator introduced in Olesen et al. (2022).
 
 Options for controlling the sliding spatial window kernel
@@ -57,14 +62,23 @@ Options for controlling the sliding spatial window kernel
 
 -  **-radius_ratio value** Set the spherical kernel size as a ratio of number of voxels to number of input volumes (default: 1.0/0.85 ~= 1.18)
 
--  **-extent window** Set the patch size of the cuboid kernel; can be either a single odd integer or a comma-separated triplet of odd integers
+-  **-extent window** Set the patch size of the cuboid kernel; can be either a single integer or a comma-separated triplet of integers (see Description)
+
+-  **-subsample factor** reduce the number of PCA kernels relative to the number of image voxels; can provide either an integer subsampling factor, or a comma-separated list of three factors;default: 2
+
+Options for phase demodulation of complex data
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+-  **-nodemod** disable phase demodulation
+
+-  **-demod_axes axes** comma-separated list of axis indices along which FFT can be applied for phase demodulation
 
 Options that affect reconstruction of the output image series
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 -  **-mask image** Only denoise voxels within the specified binary brain mask image.
 
--  **-filter choice** Modulate how component contributions are filtered based on the cumulative eigenvalues relative to the noise level; options are: truncate,frobenius; default: frobenius (Optimal Shrinkage based on minimisation of the Frobenius norm)
+-  **-filter choice** Modulate how component contributions are filtered based on the cumulative eigenvalues relative to the noise level; options are: optshrink,optthresh,truncate; default: optshrink (Optimal Shrinkage based on minimisation of the Frobenius norm)
 
 -  **-aggregator choice** Select how the outcomes of multiple PCA outcomes centred at different voxels contribute to the reconstructed DWI signal in each voxel; options are: exclusive,gaussian,invl0,rank,uniform; default: Gaussian
 
@@ -73,7 +87,7 @@ Options for exporting additional data regarding PCA behaviour
 
 -  **-noise_out image** The output noise map, i.e., the estimated noise level 'sigma' in the data. Note that on complex input data, this will be the total noise level across real and imaginary channels, so a scale factor sqrt(2) applies.
 
--  **-rank_input image** The signal rank estimated for the denoising patch centred at each input image voxel
+-  **-rank_input image** The signal rank estimated for each denoising patch
 
 -  **-rank_output image** An estimated rank for the output image data, accounting for multi-patch aggregation
 
@@ -84,6 +98,8 @@ Options for debugging the operation of sliding window kernels
 
 -  **-voxelcount image** The number of voxels that contributed to the PCA for processing of each voxel
 
+-  **-patchcount image** The number of unique patches to which an image voxel contributes
+
 -  **-sum_aggregation image** The sum of aggregation weights of those patches contributing to each output voxel
 
 -  **-sum_optshrink image** the sum of eigenvector weights computed for the denoising patch centred at each voxel as a result of performing optimal shrinkage
@@ -120,6 +136,8 @@ Cordero-Grande, L.; Christiaens, D.; Hutter, J.; Price, A.N.; Hajnal, J.V. Compl
 
 * If using anything other than -aggregation exclusive: Manjon, J.V.; Coupe, P.; Concha, L.; Buades, A.; D. Collins, D.L.; Robles, M. Diffusion Weighted Image Denoising Using Overcomplete Local PCA. PLoS ONE, 2013, 8(9), e73021
 
+* If using -estimator med or -filter optthresh: Gavish, M.; Donoho, D.L.The Optimal Hard Threshold for Singular Values is 4/sqrt(3). IEEE Transactions on Information Theory, 2014, 60(8), 5040-5053.
+
 Tournier, J.-D.; Smith, R. E.; Raffelt, D.; Tabbara, R.; Dhollander, T.; Pietsch, M.; Christiaens, D.; Jeurissen, B.; Yeh, C.-H. & Connelly, A. MRtrix3: A fast, flexible and open software framework for medical image processing and visualisation. NeuroImage, 2019, 202, 116137
 
 --------------
diff --git a/docs/reference/commands/mrfilter.rst b/docs/reference/commands/mrfilter.rst
index 5a1a7da432..fdfca7d44e 100644
--- a/docs/reference/commands/mrfilter.rst
+++ b/docs/reference/commands/mrfilter.rst
@@ -22,7 +22,7 @@ Usage
 Description
 -----------
 
-The available filters are: fft, gradient, median, smooth, normalise, zclean.
+The available filters are: demodulate, fft, gradient, median, smooth, normalise, zclean.
 
 Each filter has its own unique set of optional parameters.
 
@@ -31,11 +31,14 @@ For 4D images, each 3D volume is processed independently.
 Options
 -------
 
-Options for FFT filter
-^^^^^^^^^^^^^^^^^^^^^^
+Options applicable to both demodulate and FFT filters
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 -  **-axes list** the axes along which to apply the Fourier Transform. By default, the transform is applied along the three spatial axes. Provide as a comma-separate list of axis indices.
 
+Options for FFT filter
+^^^^^^^^^^^^^^^^^^^^^^
+
 -  **-inverse** apply the inverse FFT
 
 -  **-magnitude** output a magnitude image rather than a complex-valued image
diff --git a/src/denoise/demodulate.cpp b/src/denoise/demodulate.cpp
new file mode 100644
index 0000000000..8719a391d7
--- /dev/null
+++ b/src/denoise/demodulate.cpp
@@ -0,0 +1,88 @@
+/* Copyright (c) 2008-2024 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#include "denoise/demodulate.h"
+
+#include "app.h"
+#include "axes.h"
+
+using namespace MR::App;
+
+namespace MR::Denoise {
+
+const char *const demodulation_description = "If the input data are of complex type, "
+                                             "then a linear phase term will be removed from each k-space prior to PCA. "
+                                             "In the absence of metadata indicating otherwise, "
+                                             "it is inferred that the first two axes correspond to acquired slices, "
+                                             "and different slices / volumes will be demodulated individually; "
+                                             "this behaviour can be modified using the -demod_axes option.";
+
+const OptionGroup demodulation_options =
+    OptionGroup("Options for phase demodulation of complex data") +
+    Option("nodemod", "disable phase demodulation")
+    // TODO Consider option to disable the remodulation of the output denoised series;
+    //   would need to turn this into a function call,
+    //   as that option would need to be omitted from dwi2noise
+    // Perhaps -nodemod, -noremod could be combined into a type_choice()?
+    // This wouldn't be able to also cover the future prospect of linear vs. non-linear phase demodulation;
+    //   maybe input phase demodulation being none / linear / nonlinear would be the better type_choice()?
+    +
+    Option("demod_axes", "comma-separated list of axis indices along which FFT can be applied for phase demodulation") +
+    Argument("axes").type_sequence_int();
+
+std::vector<size_t> get_demodulation_axes(const Header &H) {
+  const bool complex = H.datatype().is_complex();
+  auto opt = App::get_options("nodemod");
+  if (!opt.empty()) {
+    if (!App::get_options("demod_axes").empty())
+      throw Exception("Options -nodemod and -demod_axes are mutually exclusive");
+    return std::vector<size_t>();
+  }
+  opt = App::get_options("demod_axes");
+  if (opt.empty()) {
+    if (complex) {
+      auto slice_encoding_it = H.keyval().find("SliceEncodingDirection");
+      if (slice_encoding_it == H.keyval().end()) {
+        INFO("No header information on slice encoding; assuming first two axes are within-slice");
+        return {0, 1};
+      } else {
+        auto dir = Axes::id2dir(slice_encoding_it->second);
+        std::vector<size_t> result;
+        for (size_t axis = 0; axis != 3; ++axis) {
+          if (!dir[axis])
+            result.push_back(axis);
+        }
+        INFO("For header SliceEncodingDirection=\"" + slice_encoding_it->second +
+             "\", "
+             "chose demodulation axes: " +
+             join(result, ","));
+        return result;
+      }
+    }
+  } else {
+    if (!complex)
+      throw Exception("Cannot perform phase demodulation on magnitude input image");
+    auto result = parse_ints<size_t>(opt[0][0]);
+    for (auto axis : result) {
+      if (axis > 2)
+        throw Exception("Phase demodulation implementation not yet robust to non-spatial axes");
+    }
+    return result;
+  }
+  return std::vector<size_t>();
+}
+
+} // namespace MR::Denoise
diff --git a/src/denoise/demodulate.h b/src/denoise/demodulate.h
new file mode 100644
index 0000000000..d55cc8078f
--- /dev/null
+++ b/src/denoise/demodulate.h
@@ -0,0 +1,32 @@
+/* Copyright (c) 2008-2024 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#pragma once
+
+#include <vector>
+
+#include "app.h"
+#include "header.h"
+
+namespace MR::Denoise {
+
+extern const char *const demodulation_description;
+
+extern const App::OptionGroup demodulation_options;
+
+std::vector<size_t> get_demodulation_axes(const Header &);
+
+} // namespace MR::Denoise
diff --git a/src/denoise/denoise.cpp b/src/denoise/denoise.cpp
index 9f86a24267..2b00d165ee 100644
--- a/src/denoise/denoise.cpp
+++ b/src/denoise/denoise.cpp
@@ -16,6 +16,8 @@
 
 #include "denoise/denoise.h"
 
+#include "axes.h"
+
 namespace MR::Denoise {
 
 using namespace App;
diff --git a/src/denoise/denoise.h b/src/denoise/denoise.h
index b35fa7f019..ee17664d34 100644
--- a/src/denoise/denoise.h
+++ b/src/denoise/denoise.h
@@ -17,8 +17,10 @@
 #pragma once
 
 #include <Eigen/Dense>
+#include <vector>
 
 #include "app.h"
+#include "header.h"
 
 namespace MR::Denoise {
 
@@ -28,8 +30,8 @@ using vector_type = Eigen::Array<double, Eigen::Dynamic, 1>;
 const std::vector<std::string> dtypes = {"float32", "float64"};
 extern const App::Option datatype_option;
 
-const std::vector<std::string> filters = {"truncate", "frobenius"};
-enum class filter_type { TRUNCATE, FROBENIUS };
+const std::vector<std::string> filters = {"optshrink", "optthresh", "truncate"};
+enum class filter_type { OPTSHRINK, OPTTHRESH, TRUNCATE };
 
 const std::vector<std::string> aggregators = {"exclusive", "gaussian", "invl0", "rank", "uniform"};
 enum class aggregator_type { EXCLUSIVE, GAUSSIAN, INVL0, RANK, UNIFORM };
diff --git a/src/denoise/estimate.cpp b/src/denoise/estimate.cpp
index a7306705e5..e8f7152cae 100644
--- a/src/denoise/estimate.cpp
+++ b/src/denoise/estimate.cpp
@@ -100,7 +100,6 @@ template <typename F> void Estimate<F>::operator()(Image<F> &dwi) {
 
   // Marchenko-Pastur optimal threshold determination
   threshold = (*estimator)(s, m, n);
-  const ssize_t in_rank = r - threshold.cutoff_p;
 
   // Store additional output maps if requested
   auto ss_index = subsample->in2ss({dwi.index(0), dwi.index(1), dwi.index(2)});
@@ -110,7 +109,7 @@ template <typename F> void Estimate<F>::operator()(Image<F> &dwi) {
   }
   if (exports.rank_input.valid()) {
     assign_pos_of(ss_index).to(exports.rank_input);
-    exports.rank_input.value() = in_rank;
+    exports.rank_input.value() = r - threshold.cutoff_p;
   }
   if (exports.max_dist.valid()) {
     assign_pos_of(ss_index).to(exports.max_dist);
diff --git a/src/denoise/estimator/estimator.cpp b/src/denoise/estimator/estimator.cpp
index 51efed00c7..911c3379d5 100644
--- a/src/denoise/estimator/estimator.cpp
+++ b/src/denoise/estimator/estimator.cpp
@@ -18,6 +18,7 @@
 
 #include "denoise/estimator/base.h"
 #include "denoise/estimator/exp.h"
+#include "denoise/estimator/med.h"
 #include "denoise/estimator/mrm2022.h"
 
 namespace MR::Denoise::Estimator {
@@ -30,6 +31,7 @@ const Option option = Option("estimator",
                              " either: \n"
                              "* Exp1: the original estimator used in Veraart et al. (2016); \n"
                              "* Exp2: the improved estimator introduced in Cordero-Grande et al. (2019); \n"
+                             "* Med: estimate based on the median eigenvalue as in Gavish and Donohue (2014); \n"
                              "* MRM2022: the alternative estimator introduced in Olesen et al. (2022).") +
                       Argument("algorithm").type_choice(estimators);
 
@@ -42,6 +44,9 @@ std::shared_ptr<Base> make_estimator() {
   case estimator_type::EXP2:
     return std::make_shared<Exp<2>>();
     break;
+  case estimator_type::MED:
+    return std::make_shared<Med>();
+    break;
   case estimator_type::MRM2022:
     return std::make_shared<MRM2022>();
     break;
diff --git a/src/denoise/estimator/estimator.h b/src/denoise/estimator/estimator.h
index f0ec397949..620db5f8bb 100644
--- a/src/denoise/estimator/estimator.h
+++ b/src/denoise/estimator/estimator.h
@@ -27,8 +27,8 @@ namespace MR::Denoise::Estimator {
 class Base;
 
 extern const App::Option option;
-const std::vector<std::string> estimators = {"exp1", "exp2", "mrm2022"};
-enum class estimator_type { EXP1, EXP2, MRM2022 };
+const std::vector<std::string> estimators = {"exp1", "exp2", "med", "mrm2022"};
+enum class estimator_type { EXP1, EXP2, MED, MRM2022 };
 std::shared_ptr<Base> make_estimator();
 
 } // namespace MR::Denoise::Estimator
diff --git a/src/denoise/estimator/med.h b/src/denoise/estimator/med.h
new file mode 100644
index 0000000000..600a9112ff
--- /dev/null
+++ b/src/denoise/estimator/med.h
@@ -0,0 +1,56 @@
+/* Copyright (c) 2008-2024 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#pragma once
+
+#include <limits>
+
+#include "denoise/estimator/base.h"
+#include "denoise/estimator/result.h"
+#include "math/math.h"
+#include "math/median.h"
+
+namespace MR::Denoise::Estimator {
+
+class Med : public Base {
+public:
+  Med() = default;
+  Result operator()(const eigenvalues_type &s, const ssize_t m, const ssize_t n) const final {
+    Result result;
+    const ssize_t beta = double(std::min(m, n)) / double(std::max(m, n));
+    // Eigenvalues should already be sorted;
+    //   no need to execute a sort for median calculation
+    const double ymed = s.size() & 1 ? s[s.size() / 2] : (0.5 * (s[s.size() / 2 - 1] + s[s.size() / 2]));
+    result.sigma2 = Math::pow2(ymed) / (std::max(m, n) * mu(beta));
+    return result;
+  }
+
+protected:
+  // Coefficients as provided in Gavish and Donohue 2014
+  // double omega(const double beta) const {
+  //   const double betasq = Math::pow2(beta);
+  //   return (0.56*beta*betasq - 0.95*betasq + 1.82*beta + 1.43);
+  // }
+  // Median of Marcenko-Pastur distribution
+  // Third-order polynomial fit to data generated using Matlab code supplementary to Gavish and Donohue 2014
+  double mu(const double beta) const {
+    const double betasq = Math::pow2(beta);
+    return ((-0.005882794526340723 * betasq * beta) - (0.007508551496715836 * betasq) - (0.3338169644754149 * beta) +
+            1.0);
+  }
+};
+
+} // namespace MR::Denoise::Estimator
diff --git a/src/denoise/estimator/result.h b/src/denoise/estimator/result.h
index 94b511301f..599d2327d4 100644
--- a/src/denoise/estimator/result.h
+++ b/src/denoise/estimator/result.h
@@ -20,7 +20,7 @@ namespace MR::Denoise::Estimator {
 
 class Result {
 public:
-  Result() : cutoff_p(0), sigma2(0.0) {}
+  Result() : cutoff_p(-1), sigma2(std::numeric_limits<double>::signaling_NaN()) {}
   ssize_t cutoff_p;
   double sigma2;
 };
diff --git a/src/denoise/recon.cpp b/src/denoise/recon.cpp
index de48013a05..3f55b38e21 100644
--- a/src/denoise/recon.cpp
+++ b/src/denoise/recon.cpp
@@ -50,6 +50,7 @@ template <typename F> void Recon<F>::operator()(Image<F> &dwi, Image<F> &out) {
   const ssize_t n = Estimate<F>::neighbourhood.voxels.size();
   const ssize_t r = std::min(Estimate<F>::m, n);
   const ssize_t q = std::max(Estimate<F>::m, n);
+  const double beta = double(r) / double(q);
   const ssize_t in_rank = r - Estimate<F>::threshold.cutoff_p;
 
   if (r > w.size())
@@ -65,19 +66,14 @@ template <typename F> void Recon<F>::operator()(Image<F> &dwi, Image<F> &out) {
   double sum_weights = 0.0;
   ssize_t out_rank = 0;
   switch (filter) {
-  case filter_type::TRUNCATE:
-    out_rank = in_rank;
-    w.head(Estimate<F>::threshold.cutoff_p).setZero();
-    w.segment(Estimate<F>::threshold.cutoff_p, in_rank).setOnes();
-    sum_weights = double(out_rank);
-    break;
-  case filter_type::FROBENIUS: {
+  case filter_type::OPTSHRINK: {
     if (Estimate<F>::threshold.sigma2 == 0.0) {
       w.head(r).setOnes();
       out_rank = r;
       sum_weights = double(r);
     } else {
-      const double beta = r / q;
+      // TODO I think this is wrong;
+      //   should be just based on a ratio of the eigenvalue to sigma?
       const double transition = 1.0 + std::sqrt(beta);
       double clam = 0.0;
       for (ssize_t i = 0; i != r; ++i) {
@@ -94,6 +90,29 @@ template <typename F> void Recon<F>::operator()(Image<F> &dwi, Image<F> &out) {
       }
     }
   } break;
+  case filter_type::OPTTHRESH: {
+    const std::map<double, double>::const_iterator it = beta2lambdastar.find(beta);
+    const double lambda_star =
+        it == beta2lambdastar.end()
+            ? sqrt(2.0 * (beta + 1.0) + ((8.0 * beta) / (beta + 1.0 + std::sqrt(Math::pow2(beta) + 14.0 * beta + 1.0))))
+            : it->second;
+    const double tau_star = lambda_star * std::sqrt(q) * std::sqrt(Estimate<F>::threshold.sigma2);
+    for (ssize_t i = 0; i != r; ++i) {
+      if (Estimate<F>::s[i] >= tau_star) {
+        w[i] = 1.0;
+        ++out_rank;
+      } else {
+        w[i] = 0.0;
+      }
+    }
+    sum_weights = out_rank;
+  } break;
+  case filter_type::TRUNCATE:
+    out_rank = in_rank;
+    w.head(Estimate<F>::threshold.cutoff_p).setZero();
+    w.segment(Estimate<F>::threshold.cutoff_p, in_rank).setOnes();
+    sum_weights = double(out_rank);
+    break;
   default:
     assert(false);
   }
diff --git a/src/denoise/subsample.cpp b/src/denoise/subsample.cpp
index ccac67c712..d4897a66bd 100644
--- a/src/denoise/subsample.cpp
+++ b/src/denoise/subsample.cpp
@@ -22,7 +22,7 @@ const App::Option subsample_option =
     App::Option("subsample",
                 "reduce the number of PCA kernels relative to the number of image voxels; "
                 "can provide either an integer subsampling factor, "
-                "or a comma-separated list of three factors;"
+                "or a comma-separated list of three factors; "
                 "default: 2") +
     App::Argument("factor").type_integer(1);
 

From 5630eccf8c44d9bf71b82d4139175dac60055ca9 Mon Sep 17 00:00:00 2001
From: Robert Smith <robert.smith@florey.edu.au>
Date: Sat, 30 Nov 2024 11:31:50 +1100
Subject: [PATCH 22/34] dwi*noise: Various fixes / changes

- Estimators now explicitly yield an estimate of the upper end of the noise distribution in addition to the estimated noise level.
- Fix to -estimator med (had gross incorrect scaling).
- Fix centering of searchlight for even subsample factors.
- Fix alignment of voxel grids between input image and output images that are defined with respect to subsampling.
- Fix creation of output data where estimator fails completely.
- Revert optimal shrinkage to use the upper bound of the MP distribution as the unity reference rather than the noise level.
- Fix optimal thresholding failing to construct the local dictionary mapping from aspect ratio to lambda_star parameter to reduce redundant calculations.
---
 cmd/dwi2noise.cpp                  |  4 +-
 src/denoise/estimator/exp.h        |  1 +
 src/denoise/estimator/med.h        | 16 ++++--
 src/denoise/estimator/mrm2022.h    |  1 +
 src/denoise/estimator/result.h     |  6 ++-
 src/denoise/kernel/sphere_base.cpp |  6 +--
 src/denoise/recon.cpp              | 79 ++++++++++++++++--------------
 src/denoise/recon.h                |  2 +
 src/denoise/subsample.cpp          | 12 ++---
 9 files changed, 74 insertions(+), 53 deletions(-)

diff --git a/cmd/dwi2noise.cpp b/cmd/dwi2noise.cpp
index e22284bc06..f2915c9e8a 100644
--- a/cmd/dwi2noise.cpp
+++ b/cmd/dwi2noise.cpp
@@ -116,10 +116,10 @@ void usage() {
     + Argument("image").type_image_out()
   + OptionGroup("Options for debugging the operation of sliding window kernels")
   + Option("max_dist",
-           "The maximum distance between a voxel and another voxel that was included in the local denoising patch")
+           "The maximum distance between the centre of the patch and a voxel that was included within that patch")
     + Argument("image").type_image_out()
   + Option("voxelcount",
-           "The number of voxels that contributed to the PCA for processing of each voxel")
+           "The number of voxels that contributed to the PCA for processing of each patch")
     + Argument("image").type_image_out()
   + Option("patchcount",
            "The number of unique patches to which an image voxel contributes")
diff --git a/src/denoise/estimator/exp.h b/src/denoise/estimator/exp.h
index 4931fe2b01..e1becedf0b 100644
--- a/src/denoise/estimator/exp.h
+++ b/src/denoise/estimator/exp.h
@@ -53,6 +53,7 @@ template <ssize_t version> class Exp : public Base {
       if (sigsq2 < sigsq1) {
         result.sigma2 = sigsq1;
         result.cutoff_p = p + 1;
+        result.lamplus = lam;
       }
     }
     return result;
diff --git a/src/denoise/estimator/med.h b/src/denoise/estimator/med.h
index 600a9112ff..14ea428d32 100644
--- a/src/denoise/estimator/med.h
+++ b/src/denoise/estimator/med.h
@@ -30,11 +30,17 @@ class Med : public Base {
   Med() = default;
   Result operator()(const eigenvalues_type &s, const ssize_t m, const ssize_t n) const final {
     Result result;
-    const ssize_t beta = double(std::min(m, n)) / double(std::max(m, n));
+    const ssize_t r = std::min(m, n);
+    const ssize_t q = std::max(m, n);
+    const ssize_t beta = double(r) / double(q);
     // Eigenvalues should already be sorted;
     //   no need to execute a sort for median calculation
     const double ymed = s.size() & 1 ? s[s.size() / 2] : (0.5 * (s[s.size() / 2 - 1] + s[s.size() / 2]));
-    result.sigma2 = Math::pow2(ymed) / (std::max(m, n) * mu(beta));
+    result.lamplus = ymed / (q * mu(beta));
+    // Mechanism intrinsically assumes half rank
+    result.cutoff_p = r / 2;
+    // Calculate noise level based on MP distribution
+    result.sigma2 = 2.0 * s.head(s.size() / 2).sum() / (q * s.size());
     return result;
   }
 
@@ -48,8 +54,10 @@ class Med : public Base {
   // Third-order polynomial fit to data generated using Matlab code supplementary to Gavish and Donohue 2014
   double mu(const double beta) const {
     const double betasq = Math::pow2(beta);
-    return ((-0.005882794526340723 * betasq * beta) - (0.007508551496715836 * betasq) - (0.3338169644754149 * beta) +
-            1.0);
+    return ((-0.005882794526340723 * betasq * beta) //
+            - (0.007508551496715836 * betasq)       //
+            - (0.3338169644754149 * beta)           //
+            + 1.0);                                 //
   }
 };
 
diff --git a/src/denoise/estimator/mrm2022.h b/src/denoise/estimator/mrm2022.h
index 1665ccb90c..4d48d836ec 100644
--- a/src/denoise/estimator/mrm2022.h
+++ b/src/denoise/estimator/mrm2022.h
@@ -50,6 +50,7 @@ class MRM2022 : public Base {
       lamplusprev = sigmasq * sigmasq_to_lamplus;
       result.cutoff_p = i;
       result.sigma2 = sigmasq;
+      result.lamplus = lamplusprev;
     }
     return result;
   }
diff --git a/src/denoise/estimator/result.h b/src/denoise/estimator/result.h
index 599d2327d4..b74b362119 100644
--- a/src/denoise/estimator/result.h
+++ b/src/denoise/estimator/result.h
@@ -20,9 +20,13 @@ namespace MR::Denoise::Estimator {
 
 class Result {
 public:
-  Result() : cutoff_p(-1), sigma2(std::numeric_limits<double>::signaling_NaN()) {}
+  Result()
+      : cutoff_p(-1),
+        sigma2(std::numeric_limits<double>::signaling_NaN()),
+        lamplus(std::numeric_limits<double>::signaling_NaN()) {}
   ssize_t cutoff_p;
   double sigma2;
+  double lamplus;
 };
 
 } // namespace MR::Denoise::Estimator
diff --git a/src/denoise/kernel/sphere_base.cpp b/src/denoise/kernel/sphere_base.cpp
index d441f2e975..666b090220 100644
--- a/src/denoise/kernel/sphere_base.cpp
+++ b/src/denoise/kernel/sphere_base.cpp
@@ -46,9 +46,9 @@ SphereBase::Shared::Shared(const Header &voxel_grid,
     for (offset[1] = bounding_box(1, 0); offset[1] <= bounding_box(1, 1); ++offset[1]) {
       for (offset[0] = bounding_box(0, 0); offset[0] <= bounding_box(0, 1); ++offset[0]) {
         const default_type squared_distance =
-            Math::pow2((offset[0] + halfvoxel_offsets[0]) * voxel_grid.spacing(0))    //
-            + Math::pow2((offset[1] + halfvoxel_offsets[1]) * voxel_grid.spacing(1))  //
-            + Math::pow2((offset[2] + halfvoxel_offsets[2]) * voxel_grid.spacing(2)); //
+            Math::pow2((offset[0] - halfvoxel_offsets[0]) * voxel_grid.spacing(0))    //
+            + Math::pow2((offset[1] - halfvoxel_offsets[1]) * voxel_grid.spacing(1))  //
+            + Math::pow2((offset[2] - halfvoxel_offsets[2]) * voxel_grid.spacing(2)); //
         if (squared_distance <= max_radius_sq)
           data.emplace_back(Offset(offset, squared_distance));
       }
diff --git a/src/denoise/recon.cpp b/src/denoise/recon.cpp
index 3f55b38e21..325e11d15d 100644
--- a/src/denoise/recon.cpp
+++ b/src/denoise/recon.cpp
@@ -65,56 +65,61 @@ template <typename F> void Recon<F>::operator()(Image<F> &dwi, Image<F> &out) {
   // Generate weights vector
   double sum_weights = 0.0;
   ssize_t out_rank = 0;
-  switch (filter) {
-  case filter_type::OPTSHRINK: {
-    if (Estimate<F>::threshold.sigma2 == 0.0) {
-      w.head(r).setOnes();
-      out_rank = r;
-      sum_weights = double(r);
-    } else {
-      // TODO I think this is wrong;
-      //   should be just based on a ratio of the eigenvalue to sigma?
+  if (Estimate<F>::threshold.sigma2 == 0.0 || !std::isfinite(Estimate<F>::threshold.sigma2)) {
+    w.head(r).setOnes();
+    out_rank = r;
+    sum_weights = double(r);
+  } else {
+    switch (filter) {
+    case filter_type::OPTSHRINK: {
       const double transition = 1.0 + std::sqrt(beta);
-      double clam = 0.0;
       for (ssize_t i = 0; i != r; ++i) {
         const double lam = std::max(Estimate<F>::s[i], 0.0) / q;
-        clam += lam;
-        const double y = clam / (Estimate<F>::threshold.sigma2 * (i + 1));
+        // TODO Should this be based on the noise level,
+        //   or on the estimated upper bound of the MP distribution?
+        // If based on upper bound,
+        //   there will be an issue with importing this information from a pre-estimated noise map
+        const double y = lam / Estimate<F>::threshold.lamplus;
         double nu = 0.0;
         if (y > transition) {
           nu = std::sqrt(Math::pow2(Math::pow2(y) - beta - 1.0) - (4.0 * beta)) / y;
           ++out_rank;
         }
-        w[i] = clam > 0.0 ? (nu / y) : 0.0;
+        w[i] = lam > 0.0 ? (nu / y) : 0.0;
+        assert(w[i] >= 0.0 && w[i] <= 1.0);
         sum_weights += w[i];
       }
-    }
-  } break;
-  case filter_type::OPTTHRESH: {
-    const std::map<double, double>::const_iterator it = beta2lambdastar.find(beta);
-    const double lambda_star =
-        it == beta2lambdastar.end()
-            ? sqrt(2.0 * (beta + 1.0) + ((8.0 * beta) / (beta + 1.0 + std::sqrt(Math::pow2(beta) + 14.0 * beta + 1.0))))
-            : it->second;
-    const double tau_star = lambda_star * std::sqrt(q) * std::sqrt(Estimate<F>::threshold.sigma2);
-    for (ssize_t i = 0; i != r; ++i) {
-      if (Estimate<F>::s[i] >= tau_star) {
-        w[i] = 1.0;
-        ++out_rank;
+    } break;
+    case filter_type::OPTTHRESH: {
+      const std::map<double, double>::const_iterator it = beta2lambdastar.find(beta);
+      double lambda_star = 0.0;
+      if (it == beta2lambdastar.end()) {
+        lambda_star =
+            sqrt(2.0 * (beta + 1.0) + ((8.0 * beta) / (beta + 1.0 + std::sqrt(Math::pow2(beta) + 14.0 * beta + 1.0))));
+        beta2lambdastar[beta] = lambda_star;
       } else {
-        w[i] = 0.0;
+        lambda_star = it->second;
       }
+      const double tau_star_sq = Math::pow2(lambda_star) * q * Estimate<F>::threshold.sigma2;
+      for (ssize_t i = 0; i != r; ++i) {
+        if (Estimate<F>::s[i] >= tau_star_sq) {
+          w[i] = 1.0;
+          ++out_rank;
+        } else {
+          w[i] = 0.0;
+        }
+      }
+      sum_weights = out_rank;
+    } break;
+    case filter_type::TRUNCATE:
+      out_rank = in_rank;
+      w.head(Estimate<F>::threshold.cutoff_p).setZero();
+      w.segment(Estimate<F>::threshold.cutoff_p, in_rank).setOnes();
+      sum_weights = double(out_rank);
+      break;
+    default:
+      assert(false);
     }
-    sum_weights = out_rank;
-  } break;
-  case filter_type::TRUNCATE:
-    out_rank = in_rank;
-    w.head(Estimate<F>::threshold.cutoff_p).setZero();
-    w.segment(Estimate<F>::threshold.cutoff_p, in_rank).setOnes();
-    sum_weights = double(out_rank);
-    break;
-  default:
-    assert(false);
   }
   assert(w.head(r).allFinite());
   assert(std::isfinite(sum_weights));
diff --git a/src/denoise/recon.h b/src/denoise/recon.h
index fd48c7ed97..a196667774 100644
--- a/src/denoise/recon.h
+++ b/src/denoise/recon.h
@@ -17,6 +17,7 @@
 #pragma once
 
 #include <limits>
+#include <map>
 #include <memory>
 
 #include <Eigen/Dense>
@@ -53,6 +54,7 @@ template <typename F> class Recon : public Estimate<F> {
   // Reusable memory
   vector_type w;
   typename Estimate<F>::MatrixType Xr;
+  std::map<double, double> beta2lambdastar;
 };
 
 template class Recon<float>;
diff --git a/src/denoise/subsample.cpp b/src/denoise/subsample.cpp
index d4897a66bd..92f5c6b8b7 100644
--- a/src/denoise/subsample.cpp
+++ b/src/denoise/subsample.cpp
@@ -92,16 +92,16 @@ Header Subsample::make_subsample_header() const {
   H.reset_intensity_scaling();
   H.datatype() = DataType::Float32;
   H.datatype().set_byte_order_native();
+  std::array<double, 3> halfvoxel_offsets;
   for (ssize_t axis = 0; axis != 3; ++axis) {
     H.size(axis) = size[axis];
     H.spacing(axis) *= factors[axis];
+    halfvoxel_offsets[axis] = factors[axis] & 1 ? 0.0 : 0.5;
   }
-  // Need to move the transform origin from voxel [0,0,0] in the input image
-  //   to voxel [0,0,0] in the subsampled voxel grid;
-  //   this is just a voxel2scanner transformation of "origin"
-  H.transform().translation() = H.transform() * Eigen::Matrix<default_type, 3, 1>({origin[0] * H_in.spacing(0),   //
-                                                                                   origin[1] * H_in.spacing(1),   //
-                                                                                   origin[2] * H_in.spacing(2)}); //
+  H.transform().translation() =
+      H_in.transform() * Eigen::Matrix<default_type, 3, 1>({(origin[0] + halfvoxel_offsets[0]) * H_in.spacing(0),   //
+                                                            (origin[1] + halfvoxel_offsets[1]) * H_in.spacing(1),   //
+                                                            (origin[2] + halfvoxel_offsets[2]) * H_in.spacing(2)}); //
   return H;
 }
 

From bc33d4d61348d9357d72114896b911b683cf0602 Mon Sep 17 00:00:00 2001
From: Robert Smith <robert.smith@florey.edu.au>
Date: Sun, 1 Dec 2024 23:00:24 +1100
Subject: [PATCH 23/34] dwidenoise: Import pre-estimated noise level map

---
 cmd/dwidenoise.cpp                  |  3 +++
 src/denoise/estimate.cpp            | 15 ++++++++++-----
 src/denoise/estimator/base.h        |  5 ++++-
 src/denoise/estimator/estimator.cpp | 15 ++++++++++++---
 src/denoise/estimator/estimator.h   |  4 ++--
 src/denoise/estimator/exp.h         |  5 ++++-
 src/denoise/estimator/med.h         |  5 ++++-
 src/denoise/estimator/mrm2022.h     |  5 ++++-
 src/denoise/subsample.cpp           |  8 ++++++++
 src/denoise/subsample.h             |  5 ++++-
 10 files changed, 55 insertions(+), 15 deletions(-)

diff --git a/cmd/dwidenoise.cpp b/cmd/dwidenoise.cpp
index cffcb62dd1..e9ee5aeeca 100644
--- a/cmd/dwidenoise.cpp
+++ b/cmd/dwidenoise.cpp
@@ -146,6 +146,9 @@ void usage() {
   + Kernel::options
   + subsample_option
   + demodulation_options
+  + Option("noise_in",
+           "import a pre-estimated noise level map rather than estimating this level during denoising")
+    + Argument("image").type_image_in()
 
   + OptionGroup("Options that affect reconstruction of the output image series")
   // TODO Separate masks for voxels to contribute to patches vs. voxels for which to perform denoising
diff --git a/src/denoise/estimate.cpp b/src/denoise/estimate.cpp
index e8f7152cae..d83e519331 100644
--- a/src/denoise/estimate.cpp
+++ b/src/denoise/estimate.cpp
@@ -50,18 +50,19 @@ template <typename F> void Estimate<F>::operator()(Image<F> &dwi) {
   //    In some use cases there may not be any image created that conforms to this voxel grid
   //    Have to transform the subsampled voxel index into an input image voxel index for the centre of the patch
   // Going to go with 1. for now, as for 2. may not have a suitable image over which to loop
-  if (!subsample->process(Kernel::Voxel::index_type({dwi.index(0), dwi.index(1), dwi.index(2)})))
+  Kernel::Voxel::index_type voxel({dwi.index(0), dwi.index(1), dwi.index(2)});
+  if (!subsample->process(voxel))
     return;
 
   // Process voxels in mask only
   if (mask.valid()) {
-    assign_pos_of(dwi, 0, 3).to(mask);
+    assign_pos_of(voxel).to(mask);
     if (!mask.value())
       return;
   }
 
   // Load list of voxels from which to load data
-  neighbourhood = (*kernel)({dwi.index(0), dwi.index(1), dwi.index(2)});
+  neighbourhood = (*kernel)(voxel);
   const ssize_t n = neighbourhood.voxels.size();
   const ssize_t r = std::min(m, n);
   const ssize_t q = std::max(m, n);
@@ -98,11 +99,15 @@ template <typename F> void Estimate<F>::operator()(Image<F> &dwi) {
   // eigenvalues sorted in increasing order:
   s.head(r) = eig.eigenvalues().template cast<double>();
 
+  // Centre of patch in realspace
+  //   (might be used by estimator)
+  const Eigen::Vector3d pos(subsample->patch_centre(voxel));
+
   // Marchenko-Pastur optimal threshold determination
-  threshold = (*estimator)(s, m, n);
+  threshold = (*estimator)(s, m, n, pos);
 
   // Store additional output maps if requested
-  auto ss_index = subsample->in2ss({dwi.index(0), dwi.index(1), dwi.index(2)});
+  auto ss_index = subsample->in2ss(voxel);
   if (exports.noise_out.valid()) {
     assign_pos_of(ss_index).to(exports.noise_out);
     exports.noise_out.value() = float(std::sqrt(threshold.sigma2));
diff --git a/src/denoise/estimator/base.h b/src/denoise/estimator/base.h
index 8b239e1850..a95b85d980 100644
--- a/src/denoise/estimator/base.h
+++ b/src/denoise/estimator/base.h
@@ -24,7 +24,10 @@ namespace MR::Denoise::Estimator {
 class Base {
 public:
   Base() = default;
-  virtual Result operator()(const eigenvalues_type &eigenvalues, const ssize_t m, const ssize_t n) const = 0;
+  virtual Result operator()(const eigenvalues_type &eigenvalues,
+                            const ssize_t m,
+                            const ssize_t n,
+                            const Eigen::Vector3d &pos) const = 0;
 };
 
 } // namespace MR::Denoise::Estimator
diff --git a/src/denoise/estimator/estimator.cpp b/src/denoise/estimator/estimator.cpp
index 911c3379d5..a8e6c92251 100644
--- a/src/denoise/estimator/estimator.cpp
+++ b/src/denoise/estimator/estimator.cpp
@@ -18,6 +18,7 @@
 
 #include "denoise/estimator/base.h"
 #include "denoise/estimator/exp.h"
+#include "denoise/estimator/import.h"
 #include "denoise/estimator/med.h"
 #include "denoise/estimator/mrm2022.h"
 
@@ -31,6 +32,7 @@ const Option option = Option("estimator",
                              " either: \n"
                              "* Exp1: the original estimator used in Veraart et al. (2016); \n"
                              "* Exp2: the improved estimator introduced in Cordero-Grande et al. (2019); \n"
+                             "* Import: import from a pre-estimated noise level map; \n"
                              "* Med: estimate based on the median eigenvalue as in Gavish and Donohue (2014); \n"
                              "* MRM2022: the alternative estimator introduced in Olesen et al. (2022).") +
                       Argument("algorithm").type_choice(estimators);
@@ -38,18 +40,25 @@ const Option option = Option("estimator",
 std::shared_ptr<Base> make_estimator() {
   auto opt = App::get_options("estimator");
   const estimator_type est = opt.empty() ? estimator_type::EXP2 : estimator_type((int)(opt[0][0]));
+  auto noise_in = get_options("noise_in");
+  // TODO Remove once -noise_in is used in other ways
+  if (!noise_in.empty() && est != estimator_type::IMPORT) {
+    WARN("-noise_in option has no effect unless -estimator import is specified");
+  }
   switch (est) {
   case estimator_type::EXP1:
     return std::make_shared<Exp<1>>();
   case estimator_type::EXP2:
     return std::make_shared<Exp<2>>();
-    break;
+  case estimator_type::IMPORT:
+    if (noise_in.empty())
+      throw Exception("-estimator import requires a pre-estimated noise level image "
+                      "to be provided via the -noise_in option");
+    return std::make_shared<Import>(noise_in[0][0]);
   case estimator_type::MED:
     return std::make_shared<Med>();
-    break;
   case estimator_type::MRM2022:
     return std::make_shared<MRM2022>();
-    break;
   default:
     assert(false);
   }
diff --git a/src/denoise/estimator/estimator.h b/src/denoise/estimator/estimator.h
index 620db5f8bb..10cae55f43 100644
--- a/src/denoise/estimator/estimator.h
+++ b/src/denoise/estimator/estimator.h
@@ -27,8 +27,8 @@ namespace MR::Denoise::Estimator {
 class Base;
 
 extern const App::Option option;
-const std::vector<std::string> estimators = {"exp1", "exp2", "med", "mrm2022"};
-enum class estimator_type { EXP1, EXP2, MED, MRM2022 };
+const std::vector<std::string> estimators = {"exp1", "exp2", "import", "med", "mrm2022"};
+enum class estimator_type { EXP1, EXP2, IMPORT, MED, MRM2022 };
 std::shared_ptr<Base> make_estimator();
 
 } // namespace MR::Denoise::Estimator
diff --git a/src/denoise/estimator/exp.h b/src/denoise/estimator/exp.h
index e1becedf0b..b0100c72e5 100644
--- a/src/denoise/estimator/exp.h
+++ b/src/denoise/estimator/exp.h
@@ -25,7 +25,10 @@ namespace MR::Denoise::Estimator {
 template <ssize_t version> class Exp : public Base {
 public:
   Exp() = default;
-  Result operator()(const eigenvalues_type &s, const ssize_t m, const ssize_t n) const final {
+  Result operator()(const eigenvalues_type &s,
+                    const ssize_t m,
+                    const ssize_t n,
+                    const Eigen::Vector3d & /*unused*/) const final {
     Result result;
     const ssize_t r = std::min(m, n);
     const ssize_t q = std::max(m, n);
diff --git a/src/denoise/estimator/med.h b/src/denoise/estimator/med.h
index 14ea428d32..c508fc4964 100644
--- a/src/denoise/estimator/med.h
+++ b/src/denoise/estimator/med.h
@@ -28,7 +28,10 @@ namespace MR::Denoise::Estimator {
 class Med : public Base {
 public:
   Med() = default;
-  Result operator()(const eigenvalues_type &s, const ssize_t m, const ssize_t n) const final {
+  Result operator()(const eigenvalues_type &s,
+                    const ssize_t m,
+                    const ssize_t n,
+                    const Eigen::Vector3d & /*unused*/) const final {
     Result result;
     const ssize_t r = std::min(m, n);
     const ssize_t q = std::max(m, n);
diff --git a/src/denoise/estimator/mrm2022.h b/src/denoise/estimator/mrm2022.h
index 4d48d836ec..f78f3146ff 100644
--- a/src/denoise/estimator/mrm2022.h
+++ b/src/denoise/estimator/mrm2022.h
@@ -27,7 +27,10 @@ namespace MR::Denoise::Estimator {
 class MRM2022 : public Base {
 public:
   MRM2022() = default;
-  Result operator()(const eigenvalues_type &s, const ssize_t m, const ssize_t n) const final {
+  Result operator()(const eigenvalues_type &s,
+                    const ssize_t m,
+                    const ssize_t n,
+                    const Eigen::Vector3d & /*unused*/) const final {
     Result result;
     const ssize_t mprime = std::min(m, n);
     const ssize_t nprime = std::max(m, n);
diff --git a/src/denoise/subsample.cpp b/src/denoise/subsample.cpp
index 92f5c6b8b7..32d9280fbc 100644
--- a/src/denoise/subsample.cpp
+++ b/src/denoise/subsample.cpp
@@ -105,4 +105,12 @@ Header Subsample::make_subsample_header() const {
   return H;
 }
 
+Eigen::Vector3d Subsample::patch_centre(const Kernel::Voxel::index_type &pos) const {
+  assert(process(pos));
+  const Eigen::Vector3d centre({double(pos[0]) + (factors[0] & 1 ? 0.0 : 0.5),
+                                double(pos[1]) + (factors[1] & 1 ? 0.0 : 0.5),
+                                double(pos[2]) + (factors[2] & 1 ? 0.0 : 0.5)});
+  return H_in.transform() * centre;
+}
+
 } // namespace MR::Denoise
diff --git a/src/denoise/subsample.h b/src/denoise/subsample.h
index f045138748..d5b8700cd9 100644
--- a/src/denoise/subsample.h
+++ b/src/denoise/subsample.h
@@ -39,13 +39,16 @@ class Subsample {
   // - Convert subsampled header voxel position to centre voxel in input image
   // TODO May want to move definition of Kernel::Voxel out of Kernel namespace
   bool process(const Kernel::Voxel::index_type &pos) const;
-  // TODO Rename these
   std::array<ssize_t, 3> in2ss(const Kernel::Voxel::index_type &pos) const;
   std::array<ssize_t, 3> ss2in(const Kernel::Voxel::index_type &pos) const;
   const std::array<ssize_t, 3> &get_factors() const { return factors; }
 
   static std::shared_ptr<Subsample> make(const Header &in);
 
+  // TODO From a processed input voxel,
+  //   get the realspace position of the centre of the patch
+  Eigen::Vector3d patch_centre(const Kernel::Voxel::index_type &pos) const;
+
 protected:
   const Header H_in;
   const std::array<ssize_t, 3> factors;

From 304e33b53cad410fc479c5bb2d442f851cbb1a3b Mon Sep 17 00:00:00 2001
From: Robert Smith <robert.smith@florey.edu.au>
Date: Sun, 1 Dec 2024 23:05:55 +1100
Subject: [PATCH 24/34] dwidenoise: Remove -mask option

Functionality is incompatible with latest enhancements to / capabilities of denoising.
---
 cmd/dwi2noise.cpp        | 10 ++--------
 cmd/dwidenoise.cpp       | 29 ++++++++---------------------
 src/denoise/estimate.cpp |  9 ---------
 src/denoise/estimate.h   |  2 --
 src/denoise/recon.cpp    |  3 +--
 src/denoise/recon.h      |  1 -
 6 files changed, 11 insertions(+), 43 deletions(-)

diff --git a/cmd/dwi2noise.cpp b/cmd/dwi2noise.cpp
index f2915c9e8a..e9d7a81877 100644
--- a/cmd/dwi2noise.cpp
+++ b/cmd/dwi2noise.cpp
@@ -106,10 +106,6 @@ void usage() {
   + subsample_option
   + demodulation_options
 
-
-  // TODO Implement mask option
-  // Note that behaviour of -mask for dwi2noise may be different to that of dwidenoise
-
   + OptionGroup("Options for exporting additional data regarding PCA behaviour")
   + Option("rank",
            "The signal rank estimated for the denoising patch centred at each input image voxel")
@@ -161,8 +157,7 @@ void run(Header &data,
          std::shared_ptr<Estimator::Base> estimator,
          Exports &exports) {
   auto input = data.get_image<T>().with_direct_io(3);
-  Image<bool> mask; // unused
-  Estimate<T> func(data, mask, subsample, kernel, estimator, exports);
+  Estimate<T> func(data, subsample, kernel, estimator, exports);
   ThreadedLoop("running MP-PCA noise level estimation", data, 0, 3).run(func, input);
 }
 
@@ -183,8 +178,7 @@ void run(Header &data,
     Filter::Demodulate demodulator(input, demodulation_axes);
     demodulator(input, input_demod);
   }
-  Image<bool> mask; // unused
-  Estimate<T> func(data, mask, subsample, kernel, estimator, exports);
+  Estimate<T> func(data, subsample, kernel, estimator, exports);
   ThreadedLoop("running MP-PCA noise level estimation", data, 0, 3).run(func, input_demod);
 }
 
diff --git a/cmd/dwidenoise.cpp b/cmd/dwidenoise.cpp
index e9ee5aeeca..b6e5d0a7ef 100644
--- a/cmd/dwidenoise.cpp
+++ b/cmd/dwidenoise.cpp
@@ -151,10 +151,6 @@ void usage() {
     + Argument("image").type_image_in()
 
   + OptionGroup("Options that affect reconstruction of the output image series")
-  // TODO Separate masks for voxels to contribute to patches vs. voxels for which to perform denoising
-  + Option("mask",
-           "Only denoise voxels within the specified binary brain mask image.")
-    + Argument("image").type_image_in()
   + Option("filter",
            "Modulate how component contributions are filtered "
            "based on the cumulative eigenvalues relative to the noise level; "
@@ -257,7 +253,6 @@ std::complex<double> operator/(const std::complex<double> &c, const float n) { r
 
 template <typename T>
 void run(Header &data,
-         Image<bool> &mask,
          std::shared_ptr<Subsample> subsample,
          std::shared_ptr<Kernel::Base> kernel,
          std::shared_ptr<Estimator::Base> estimator,
@@ -271,7 +266,7 @@ void run(Header &data,
   header.datatype() = DataType::from<T>();
   auto output = Image<T>::create(output_name, header);
   // run
-  Recon<T> func(data, mask, subsample, kernel, estimator, filter, aggregator, exports);
+  Recon<T> func(data, subsample, kernel, estimator, filter, aggregator, exports);
   ThreadedLoop("running MP-PCA denoising", data, 0, 3).run(func, input, output);
   // Rescale output if performing aggregation
   if (aggregator == aggregator_type::EXCLUSIVE)
@@ -288,7 +283,6 @@ void run(Header &data,
 
 template <typename T>
 void run(Header &data,
-         Image<bool> &mask,
          const std::vector<size_t> &demodulation_axes,
          std::shared_ptr<Subsample> subsample,
          std::shared_ptr<Kernel::Base> kernel,
@@ -298,7 +292,7 @@ void run(Header &data,
          const std::string &output_name,
          Exports &exports) {
   if (demodulation_axes.empty()) {
-    run<T>(data, mask, subsample, kernel, estimator, filter, aggregator, output_name, exports);
+    run<T>(data, subsample, kernel, estimator, filter, aggregator, output_name, exports);
     return;
   }
   auto input = data.get_image<T>();
@@ -316,7 +310,7 @@ void run(Header &data,
   header.datatype() = DataType::from<T>();
   auto output = Image<T>::create(output_name, header);
   // run
-  Recon<T> func(data, mask, subsample, kernel, estimator, filter, aggregator, exports);
+  Recon<T> func(data, subsample, kernel, estimator, filter, aggregator, exports);
   ThreadedLoop("running MP-PCA denoising", data, 0, 3).run(func, input_demodulated, output);
   // Re-apply phase ramps that were previously demodulated
   demodulate(output, true);
@@ -339,13 +333,6 @@ void run() {
   if (dwi.ndim() != 4 || dwi.size(3) <= 1)
     throw Exception("input image must be 4-dimensional");
 
-  Image<bool> mask;
-  auto opt = get_options("mask");
-  if (!opt.empty()) {
-    mask = Image<bool>::open(opt[0][0]);
-    check_dimensions(mask, dwi, 0, 3);
-  }
-
   auto subsample = Subsample::make(dwi);
   assert(subsample);
 
@@ -356,7 +343,7 @@ void run() {
   assert(estimator);
 
   filter_type filter = filter_type::OPTSHRINK;
-  opt = get_options("filter");
+  auto opt = get_options("filter");
   if (!opt.empty())
     filter = filter_type(int(opt[0][0]));
 
@@ -424,20 +411,20 @@ void run() {
   case 0:
     assert(demodulation_axes.empty());
     INFO("select real float32 for processing");
-    run<float>(dwi, mask, subsample, kernel, estimator, filter, aggregator, argument[1], exports);
+    run<float>(dwi, subsample, kernel, estimator, filter, aggregator, argument[1], exports);
     break;
   case 1:
     assert(demodulation_axes.empty());
     INFO("select real float64 for processing");
-    run<double>(dwi, mask, subsample, kernel, estimator, filter, aggregator, argument[1], exports);
+    run<double>(dwi, subsample, kernel, estimator, filter, aggregator, argument[1], exports);
     break;
   case 2:
     INFO("select complex float32 for processing");
-    run<cfloat>(dwi, mask, demodulation_axes, subsample, kernel, estimator, filter, aggregator, argument[1], exports);
+    run<cfloat>(dwi, demodulation_axes, subsample, kernel, estimator, filter, aggregator, argument[1], exports);
     break;
   case 3:
     INFO("select complex float64 for processing");
-    run<cdouble>(dwi, mask, demodulation_axes, subsample, kernel, estimator, filter, aggregator, argument[1], exports);
+    run<cdouble>(dwi, demodulation_axes, subsample, kernel, estimator, filter, aggregator, argument[1], exports);
     break;
   }
 }
diff --git a/src/denoise/estimate.cpp b/src/denoise/estimate.cpp
index d83e519331..a900694305 100644
--- a/src/denoise/estimate.cpp
+++ b/src/denoise/estimate.cpp
@@ -24,13 +24,11 @@ namespace MR::Denoise {
 
 template <typename F>
 Estimate<F>::Estimate(const Header &header,
-                      Image<bool> &mask,
                       std::shared_ptr<Subsample> subsample,
                       std::shared_ptr<Kernel::Base> kernel,
                       std::shared_ptr<Estimator::Base> estimator,
                       Exports &exports)
     : m(header.size(3)),
-      mask(mask),
       subsample(subsample),
       kernel(kernel),
       estimator(estimator),
@@ -54,13 +52,6 @@ template <typename F> void Estimate<F>::operator()(Image<F> &dwi) {
   if (!subsample->process(voxel))
     return;
 
-  // Process voxels in mask only
-  if (mask.valid()) {
-    assign_pos_of(voxel).to(mask);
-    if (!mask.value())
-      return;
-  }
-
   // Load list of voxels from which to load data
   neighbourhood = (*kernel)(voxel);
   const ssize_t n = neighbourhood.voxels.size();
diff --git a/src/denoise/estimate.h b/src/denoise/estimate.h
index c8f26f423c..8700ed1420 100644
--- a/src/denoise/estimate.h
+++ b/src/denoise/estimate.h
@@ -40,7 +40,6 @@ template <typename F> class Estimate {
   using MatrixType = Eigen::Matrix<F, Eigen::Dynamic, Eigen::Dynamic>;
 
   Estimate(const Header &header,
-           Image<bool> &mask,
            std::shared_ptr<Subsample> subsample,
            std::shared_ptr<Kernel::Base> kernel,
            std::shared_ptr<Estimator::Base> estimator,
@@ -52,7 +51,6 @@ template <typename F> class Estimate {
   const ssize_t m;
 
   // Denoising configuration
-  Image<bool> mask;
   std::shared_ptr<Subsample> subsample;
   std::shared_ptr<Kernel::Base> kernel;
   std::shared_ptr<Estimator::Base> estimator;
diff --git a/src/denoise/recon.cpp b/src/denoise/recon.cpp
index 325e11d15d..a3d467941f 100644
--- a/src/denoise/recon.cpp
+++ b/src/denoise/recon.cpp
@@ -22,14 +22,13 @@ namespace MR::Denoise {
 
 template <typename F>
 Recon<F>::Recon(const Header &header,
-                Image<bool> &mask,
                 std::shared_ptr<Subsample> subsample,
                 std::shared_ptr<Kernel::Base> kernel,
                 std::shared_ptr<Estimator::Base> estimator,
                 filter_type filter,
                 aggregator_type aggregator,
                 Exports &exports)
-    : Estimate<F>(header, mask, subsample, kernel, estimator, exports),
+    : Estimate<F>(header, subsample, kernel, estimator, exports),
       filter(filter),
       aggregator(aggregator),
       // FWHM = 2 x cube root of spacings between kernels
diff --git a/src/denoise/recon.h b/src/denoise/recon.h
index a196667774..6db9e40d59 100644
--- a/src/denoise/recon.h
+++ b/src/denoise/recon.h
@@ -35,7 +35,6 @@ template <typename F> class Recon : public Estimate<F> {
 
 public:
   Recon(const Header &header,
-        Image<bool> &mask,
         std::shared_ptr<Subsample> subsample,
         std::shared_ptr<Kernel::Base> kernel,
         std::shared_ptr<Estimator::Base> estimator,

From 67c901572c15e0844a399a3cd2e2a22658a6baa9 Mon Sep 17 00:00:00 2001
From: Robert Smith <robert.smith@florey.edu.au>
Date: Mon, 2 Dec 2024 15:21:40 +1100
Subject: [PATCH 25/34] dwidenoise: Fixes to optimal shrinkage / thresholding

---
 src/denoise/estimate.h |  1 -
 src/denoise/recon.cpp  | 10 +++++++---
 2 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/src/denoise/estimate.h b/src/denoise/estimate.h
index 8700ed1420..6ae7f72491 100644
--- a/src/denoise/estimate.h
+++ b/src/denoise/estimate.h
@@ -62,7 +62,6 @@ template <typename F> class Estimate {
   Eigen::SelfAdjointEigenSolver<MatrixType> eig;
   eigenvalues_type s;
   Estimator::Result threshold;
-  vector_type clam;
 
   // Export images
   Exports exports;
diff --git a/src/denoise/recon.cpp b/src/denoise/recon.cpp
index a3d467941f..4eb5df51e3 100644
--- a/src/denoise/recon.cpp
+++ b/src/denoise/recon.cpp
@@ -78,7 +78,9 @@ template <typename F> void Recon<F>::operator()(Image<F> &dwi, Image<F> &out) {
         //   or on the estimated upper bound of the MP distribution?
         // If based on upper bound,
         //   there will be an issue with importing this information from a pre-estimated noise map
-        const double y = lam / Estimate<F>::threshold.lamplus;
+        // TODO Unexpected absence of sqrt() here
+        // const double y = lam / std::sqrt(Estimate<F>::threshold.sigma2);
+        const double y = lam / Estimate<F>::threshold.sigma2;
         double nu = 0.0;
         if (y > transition) {
           nu = std::sqrt(Math::pow2(Math::pow2(y) - beta - 1.0) - (4.0 * beta)) / y;
@@ -99,9 +101,11 @@ template <typename F> void Recon<F>::operator()(Image<F> &dwi, Image<F> &out) {
       } else {
         lambda_star = it->second;
       }
-      const double tau_star_sq = Math::pow2(lambda_star) * q * Estimate<F>::threshold.sigma2;
+      const double tau_star = lambda_star * std::sqrt(q) * std::sqrt(Estimate<F>::threshold.sigma2);
+      // TODO Unexpected requisite square applied to q here
+      const double threshold = tau_star * Math::pow2(q);
       for (ssize_t i = 0; i != r; ++i) {
-        if (Estimate<F>::s[i] >= tau_star_sq) {
+        if (Estimate<F>::s[i] >= threshold) {
           w[i] = 1.0;
           ++out_rank;
         } else {

From a680743b0eaff6bc0eff7fe2efac31d2a6e80dc8 Mon Sep 17 00:00:00 2001
From: Robert Smith <robert.smith@florey.edu.au>
Date: Tue, 3 Dec 2024 22:16:49 +1100
Subject: [PATCH 26/34] dwi*noise: New option -nonstationarity

Option gives the ability to compensate for pre-estimated non-fixed noise level within each spatial patch.
---
 cmd/dwi2noise.cpp                    | 34 ++++++++-----
 cmd/dwidenoise.cpp                   | 74 +++++++++++++++++-----------
 src/denoise/estimate.cpp             | 48 ++++++++++++------
 src/denoise/estimate.h               |  6 ++-
 src/denoise/estimator/estimator.cpp  | 23 ++++-----
 src/denoise/estimator/estimator.h    |  6 +--
 src/denoise/kernel/base.h            | 22 ++++++++-
 src/denoise/kernel/cuboid.cpp        | 10 ++--
 src/denoise/kernel/cuboid.h          |  3 +-
 src/denoise/kernel/data.h            | 21 +++++++-
 src/denoise/kernel/kernel.cpp        |  6 +--
 src/denoise/kernel/sphere_base.cpp   |  8 ++-
 src/denoise/kernel/sphere_base.h     | 12 +++--
 src/denoise/kernel/sphere_radius.cpp |  2 +-
 src/denoise/kernel/sphere_radius.h   |  6 +--
 src/denoise/kernel/sphere_ratio.cpp  |  2 +-
 src/denoise/kernel/sphere_ratio.h    |  4 +-
 src/denoise/kernel/voxel.h           | 15 +++++-
 src/denoise/recon.cpp                | 41 ++++++++-------
 src/denoise/recon.h                  |  1 +
 src/denoise/subsample.cpp            |  8 ---
 src/denoise/subsample.h              |  8 ---
 22 files changed, 223 insertions(+), 137 deletions(-)

diff --git a/cmd/dwi2noise.cpp b/cmd/dwi2noise.cpp
index e9d7a81877..f64bc25333 100644
--- a/cmd/dwi2noise.cpp
+++ b/cmd/dwi2noise.cpp
@@ -105,10 +105,15 @@ void usage() {
   + Kernel::options
   + subsample_option
   + demodulation_options
+  + Option("nonstationarity",
+           "import an estimated map of noise nonstationarity; "
+           "note that this will be used for within-patch non-stationariy correction only, "
+           "the output noise level estimate will still be derived from the input data")
+    + Argument("image").type_image_in()
 
   + OptionGroup("Options for exporting additional data regarding PCA behaviour")
   + Option("rank",
-           "The signal rank estimated for the denoising patch centred at each input image voxel")
+           "The signal rank estimated for each denoising patch")
     + Argument("image").type_image_out()
   + OptionGroup("Options for debugging the operation of sliding window kernels")
   + Option("max_dist",
@@ -118,7 +123,7 @@ void usage() {
            "The number of voxels that contributed to the PCA for processing of each patch")
     + Argument("image").type_image_out()
   + Option("patchcount",
-           "The number of unique patches to which an image voxel contributes")
+           "The number of unique patches to which an input image voxel contributes")
     + Argument("image").type_image_out();
 
   COPYRIGHT =
@@ -154,10 +159,11 @@ template <typename T>
 void run(Header &data,
          std::shared_ptr<Subsample> subsample,
          std::shared_ptr<Kernel::Base> kernel,
+         Image<float> &nonstationarity_image,
          std::shared_ptr<Estimator::Base> estimator,
          Exports &exports) {
   auto input = data.get_image<T>().with_direct_io(3);
-  Estimate<T> func(data, subsample, kernel, estimator, exports);
+  Estimate<T> func(data, subsample, kernel, nonstationarity_image, estimator, exports);
   ThreadedLoop("running MP-PCA noise level estimation", data, 0, 3).run(func, input);
 }
 
@@ -166,10 +172,11 @@ void run(Header &data,
          const std::vector<size_t> &demodulation_axes,
          std::shared_ptr<Subsample> subsample,
          std::shared_ptr<Kernel::Base> kernel,
+         Image<float> &nonstationarity_image,
          std::shared_ptr<Estimator::Base> estimator,
          Exports &exports) {
   if (demodulation_axes.empty()) {
-    run<T>(data, subsample, kernel, estimator, exports);
+    run<T>(data, subsample, kernel, nonstationarity_image, estimator, exports);
     return;
   }
   auto input = data.get_image<T>();
@@ -178,7 +185,7 @@ void run(Header &data,
     Filter::Demodulate demodulator(input, demodulation_axes);
     demodulator(input, input_demod);
   }
-  Estimate<T> func(data, subsample, kernel, estimator, exports);
+  Estimate<T> func(data, subsample, kernel, nonstationarity_image, estimator, exports);
   ThreadedLoop("running MP-PCA noise level estimation", data, 0, 3).run(func, input_demod);
 }
 
@@ -188,18 +195,23 @@ void run() {
     throw Exception("input image must be 4-dimensional");
   bool complex = dwi.datatype().is_complex();
 
+  Image<float> nonstationarity_image;
+  auto opt = get_options("nonstationarity");
+  if (!opt.empty())
+    nonstationarity_image = Image<float>::open(opt[0][0]);
+
   auto subsample = Subsample::make(dwi);
   assert(subsample);
 
   auto kernel = Kernel::make_kernel(dwi, subsample->get_factors());
   assert(kernel);
 
-  auto estimator = Estimator::make_estimator();
+  auto estimator = Estimator::make_estimator(false);
   assert(estimator);
 
   Exports exports(dwi, subsample->header());
   exports.set_noise_out(argument[1]);
-  auto opt = get_options("rank");
+  opt = get_options("rank");
   if (!opt.empty())
     exports.set_rank_input(opt[0][0]);
   opt = get_options("max_dist");
@@ -221,20 +233,20 @@ void run() {
   case 0:
     assert(demodulation_axes.empty());
     INFO("select real float32 for processing");
-    run<float>(dwi, subsample, kernel, estimator, exports);
+    run<float>(dwi, subsample, kernel, nonstationarity_image, estimator, exports);
     break;
   case 1:
     assert(demodulation_axes.empty());
     INFO("select real float64 for processing");
-    run<double>(dwi, subsample, kernel, estimator, exports);
+    run<double>(dwi, subsample, kernel, nonstationarity_image, estimator, exports);
     break;
   case 2:
     INFO("select complex float32 for processing");
-    run<cfloat>(dwi, demodulation_axes, subsample, kernel, estimator, exports);
+    run<cfloat>(dwi, demodulation_axes, subsample, kernel, nonstationarity_image, estimator, exports);
     break;
   case 3:
     INFO("select complex float64 for processing");
-    run<cdouble>(dwi, demodulation_axes, subsample, kernel, estimator, exports);
+    run<cdouble>(dwi, demodulation_axes, subsample, kernel, nonstationarity_image, estimator, exports);
     break;
   }
 }
diff --git a/cmd/dwidenoise.cpp b/cmd/dwidenoise.cpp
index b6e5d0a7ef..c50c90a452 100644
--- a/cmd/dwidenoise.cpp
+++ b/cmd/dwidenoise.cpp
@@ -146,8 +146,16 @@ void usage() {
   + Kernel::options
   + subsample_option
   + demodulation_options
+  // TODO If explicitly regressing the mean prior to Casorati formation,
+  //   this should happen _before_ rescaling based on noise level
+  + Option("nonstationarity",
+           "import an estimated map of noise nonstationarity; "
+           "note that this will be used for within-patch non-stationariy correction only, "
+           "if noise level estimate is to be used for denoising also "
+           "it must be additionally provided via the -noise_in option")
+    + Argument("image").type_image_in()
   + Option("noise_in",
-           "import a pre-estimated noise level map rather than estimating this level during denoising")
+           "import a pre-estimated noise level map for noise removal rather than estimating this level from data")
     + Argument("image").type_image_in()
 
   + OptionGroup("Options that affect reconstruction of the output image series")
@@ -165,25 +173,6 @@ void usage() {
   // TODO For specifically the Gaussian aggregator,
   //   should ideally be possible to select the FWHM of the aggregator
 
-  // TODO Consider restructuring & encapsulating in classes
-  // Also bear in mind how use of subsampling could affect this
-  // TODO If using downsampling,
-  //   may be preferable to explicitly downsample the voxel grid on which to yield these data
-  // - -noise: This is a patch-wise estimate;
-  //           could however elect to make it an aggreate mean if aggregation is being performed
-  //           (or just forbid it if downsampling is performed?)
-  // - -rank: This is a patch-wise estimate (see -noise above)
-  // - -weightedrank: This is explicitly a multi-patch aggregation,
-  //   so doesn't apply to noise level estimation only;
-  //   it may be better to distinguish between "input rank" and "output rank"?
-  // - -sumweights: Explicitly a multi-patch aggregation
-  // - -max_dist: This is a property of the local patch;
-  //              this could be hidden behind a #define TBH
-  // - -voxels: This is a property of the local patch (see -max_dist above)
-  // - -aggregation_sum: This is explicitly a multi-patch aggregation
-  // TODO Consider an option group for "debugging of sliding window kernel behaviour"
-  // TODO Potential issue wherein optimal shrinkage may set to 0
-  //   some components above the determined noise level...
   + OptionGroup("Options for exporting additional data regarding PCA behaviour")
   + Option("noise_out",
            "The output noise map,"
@@ -255,6 +244,7 @@ template <typename T>
 void run(Header &data,
          std::shared_ptr<Subsample> subsample,
          std::shared_ptr<Kernel::Base> kernel,
+         Image<float> &nonstationarity_image,
          std::shared_ptr<Estimator::Base> estimator,
          filter_type filter,
          aggregator_type aggregator,
@@ -266,9 +256,9 @@ void run(Header &data,
   header.datatype() = DataType::from<T>();
   auto output = Image<T>::create(output_name, header);
   // run
-  Recon<T> func(data, subsample, kernel, estimator, filter, aggregator, exports);
+  Recon<T> func(data, subsample, kernel, nonstationarity_image, estimator, filter, aggregator, exports);
   ThreadedLoop("running MP-PCA denoising", data, 0, 3).run(func, input, output);
-  // Rescale output if performing aggregation
+  // Rescale output if aggregation was performed
   if (aggregator == aggregator_type::EXCLUSIVE)
     return;
   for (auto l_voxel = Loop(exports.sum_aggregation)(output, exports.sum_aggregation); l_voxel; ++l_voxel) {
@@ -286,13 +276,14 @@ void run(Header &data,
          const std::vector<size_t> &demodulation_axes,
          std::shared_ptr<Subsample> subsample,
          std::shared_ptr<Kernel::Base> kernel,
+         Image<float> &nonstationarity_image,
          std::shared_ptr<Estimator::Base> estimator,
          filter_type filter,
          aggregator_type aggregator,
          const std::string &output_name,
          Exports &exports) {
   if (demodulation_axes.empty()) {
-    run<T>(data, subsample, kernel, estimator, filter, aggregator, output_name, exports);
+    run<T>(data, subsample, kernel, nonstationarity_image, estimator, filter, aggregator, output_name, exports);
     return;
   }
   auto input = data.get_image<T>();
@@ -310,7 +301,7 @@ void run(Header &data,
   header.datatype() = DataType::from<T>();
   auto output = Image<T>::create(output_name, header);
   // run
-  Recon<T> func(data, subsample, kernel, estimator, filter, aggregator, exports);
+  Recon<T> func(data, subsample, kernel, nonstationarity_image, estimator, filter, aggregator, exports);
   ThreadedLoop("running MP-PCA denoising", data, 0, 3).run(func, input_demodulated, output);
   // Re-apply phase ramps that were previously demodulated
   demodulate(output, true);
@@ -339,11 +330,16 @@ void run() {
   auto kernel = Kernel::make_kernel(dwi, subsample->get_factors());
   assert(kernel);
 
-  auto estimator = Estimator::make_estimator();
+  Image<float> nonstationarity_image;
+  auto opt = get_options("nonstationarity");
+  if (!opt.empty())
+    nonstationarity_image = Image<float>::open(opt[0][0]);
+
+  auto estimator = Estimator::make_estimator(true);
   assert(estimator);
 
   filter_type filter = filter_type::OPTSHRINK;
-  auto opt = get_options("filter");
+  opt = get_options("filter");
   if (!opt.empty())
     filter = filter_type(int(opt[0][0]));
 
@@ -411,20 +407,38 @@ void run() {
   case 0:
     assert(demodulation_axes.empty());
     INFO("select real float32 for processing");
-    run<float>(dwi, subsample, kernel, estimator, filter, aggregator, argument[1], exports);
+    run<float>(dwi, subsample, kernel, nonstationarity_image, estimator, filter, aggregator, argument[1], exports);
     break;
   case 1:
     assert(demodulation_axes.empty());
     INFO("select real float64 for processing");
-    run<double>(dwi, subsample, kernel, estimator, filter, aggregator, argument[1], exports);
+    run<double>(dwi, subsample, kernel, nonstationarity_image, estimator, filter, aggregator, argument[1], exports);
     break;
   case 2:
     INFO("select complex float32 for processing");
-    run<cfloat>(dwi, demodulation_axes, subsample, kernel, estimator, filter, aggregator, argument[1], exports);
+    run<cfloat>(dwi,
+                demodulation_axes,
+                subsample,
+                kernel,
+                nonstationarity_image,
+                estimator,
+                filter,
+                aggregator,
+                argument[1],
+                exports);
     break;
   case 3:
     INFO("select complex float64 for processing");
-    run<cdouble>(dwi, demodulation_axes, subsample, kernel, estimator, filter, aggregator, argument[1], exports);
+    run<cdouble>(dwi,
+                 demodulation_axes,
+                 subsample,
+                 kernel,
+                 nonstationarity_image,
+                 estimator,
+                 filter,
+                 aggregator,
+                 argument[1],
+                 exports);
     break;
   }
 }
diff --git a/src/denoise/estimate.cpp b/src/denoise/estimate.cpp
index a900694305..0b6e7be5c1 100644
--- a/src/denoise/estimate.cpp
+++ b/src/denoise/estimate.cpp
@@ -18,6 +18,7 @@
 
 #include <limits>
 
+#include "interp/cubic.h"
 #include "math/math.h"
 
 namespace MR::Denoise {
@@ -26,12 +27,14 @@ template <typename F>
 Estimate<F>::Estimate(const Header &header,
                       std::shared_ptr<Subsample> subsample,
                       std::shared_ptr<Kernel::Base> kernel,
+                      Image<float> &nonstationarity_image,
                       std::shared_ptr<Estimator::Base> estimator,
                       Exports &exports)
     : m(header.size(3)),
       subsample(subsample),
       kernel(kernel),
       estimator(estimator),
+      nonstationarity_image(nonstationarity_image),
       X(m, kernel->estimated_size()),
       XtX(std::min(m, kernel->estimated_size()), std::min(m, kernel->estimated_size())),
       eig(std::min(m, kernel->estimated_size())),
@@ -52,9 +55,9 @@ template <typename F> void Estimate<F>::operator()(Image<F> &dwi) {
   if (!subsample->process(voxel))
     return;
 
-  // Load list of voxels from which to load data
-  neighbourhood = (*kernel)(voxel);
-  const ssize_t n = neighbourhood.voxels.size();
+  // Load list of voxels from which to import data
+  patch = (*kernel)(voxel);
+  const ssize_t n = patch.voxels.size();
   const ssize_t r = std::min(m, n);
   const ssize_t q = std::max(m, n);
 
@@ -79,7 +82,7 @@ template <typename F> void Estimate<F>::operator()(Image<F> &dwi) {
   s.fill(std::numeric_limits<default_type>::signaling_NaN());
 #endif
 
-  load_data(dwi, neighbourhood.voxels);
+  load_data(dwi);
 
   // Compute Eigendecomposition:
   if (m <= n)
@@ -90,12 +93,8 @@ template <typename F> void Estimate<F>::operator()(Image<F> &dwi) {
   // eigenvalues sorted in increasing order:
   s.head(r) = eig.eigenvalues().template cast<double>();
 
-  // Centre of patch in realspace
-  //   (might be used by estimator)
-  const Eigen::Vector3d pos(subsample->patch_centre(voxel));
-
   // Marchenko-Pastur optimal threshold determination
-  threshold = (*estimator)(s, m, n, pos);
+  threshold = (*estimator)(s, m, n, patch.centre_realspace);
 
   // Store additional output maps if requested
   auto ss_index = subsample->in2ss(voxel);
@@ -109,7 +108,7 @@ template <typename F> void Estimate<F>::operator()(Image<F> &dwi) {
   }
   if (exports.max_dist.valid()) {
     assign_pos_of(ss_index).to(exports.max_dist);
-    exports.max_dist.value() = neighbourhood.max_distance;
+    exports.max_dist.value() = patch.max_distance;
   }
   if (exports.voxelcount.valid()) {
     assign_pos_of(ss_index).to(exports.voxelcount);
@@ -117,18 +116,37 @@ template <typename F> void Estimate<F>::operator()(Image<F> &dwi) {
   }
   if (exports.patchcount.valid()) {
     std::lock_guard<std::mutex> lock(Estimate<F>::mutex);
-    for (const auto &v : neighbourhood.voxels) {
+    for (const auto &v : patch.voxels) {
       assign_pos_of(v.index).to(exports.patchcount);
       exports.patchcount.value() = exports.patchcount.value() + 1;
     }
   }
 }
 
-template <typename F> void Estimate<F>::load_data(Image<F> &image, const std::vector<Kernel::Voxel> &voxels) {
+template <typename F> void Estimate<F>::load_data(Image<F> &image) {
   const Kernel::Voxel::index_type pos({image.index(0), image.index(1), image.index(2)});
-  for (ssize_t i = 0; i != voxels.size(); ++i) {
-    assign_pos_of(voxels[i].index, 0, 3).to(image);
-    X.col(i) = image.row(3);
+  if (nonstationarity_image.valid()) {
+    assert(patch.centre_realspace.allFinite());
+    Interp::Cubic<Image<float>> interp(nonstationarity_image);
+    interp.scanner(patch.centre_realspace);
+    assert(!(!interp));
+    patch.centre_noise = interp.value();
+    for (ssize_t i = 0; i != patch.voxels.size(); ++i) {
+      interp.scanner(image.transform() * patch.voxels[i].index.cast<default_type>());
+      // TODO Trying to pull intensity information from voxels beyond the extremities of the subsampled image
+      //   may cause problems
+      assert(!(!interp));
+      const double voxel_noise = interp.value();
+      patch.voxels[i].noise_level = voxel_noise;
+      assign_pos_of(patch.voxels[i].index, 0, 3).to(image);
+      X.col(i) = image.row(3);
+      X.col(i) *= patch.centre_noise / voxel_noise;
+    }
+  } else {
+    for (ssize_t i = 0; i != patch.voxels.size(); ++i) {
+      assign_pos_of(patch.voxels[i].index, 0, 3).to(image);
+      X.col(i) = image.row(3);
+    }
   }
   assign_pos_of(pos, 0, 3).to(image);
 }
diff --git a/src/denoise/estimate.h b/src/denoise/estimate.h
index 6ae7f72491..82f0f6dcc8 100644
--- a/src/denoise/estimate.h
+++ b/src/denoise/estimate.h
@@ -42,6 +42,7 @@ template <typename F> class Estimate {
   Estimate(const Header &header,
            std::shared_ptr<Subsample> subsample,
            std::shared_ptr<Kernel::Base> kernel,
+           Image<float> &nonstationarity_image,
            std::shared_ptr<Estimator::Base> estimator,
            Exports &exports);
 
@@ -56,7 +57,8 @@ template <typename F> class Estimate {
   std::shared_ptr<Estimator::Base> estimator;
 
   // Reusable memory
-  Kernel::Data neighbourhood;
+  Kernel::Data patch;
+  Image<float> nonstationarity_image;
   MatrixType X;
   MatrixType XtX;
   Eigen::SelfAdjointEigenSolver<MatrixType> eig;
@@ -69,7 +71,7 @@ template <typename F> class Estimate {
   // Some data can only be written in a thread-safe manner
   static std::mutex mutex;
 
-  void load_data(Image<F> &image, const std::vector<Kernel::Voxel> &voxels);
+  void load_data(Image<F> &image);
 };
 
 template <typename F> std::mutex Estimate<F>::mutex;
diff --git a/src/denoise/estimator/estimator.cpp b/src/denoise/estimator/estimator.cpp
index a8e6c92251..7720bf5254 100644
--- a/src/denoise/estimator/estimator.cpp
+++ b/src/denoise/estimator/estimator.cpp
@@ -32,29 +32,26 @@ const Option option = Option("estimator",
                              " either: \n"
                              "* Exp1: the original estimator used in Veraart et al. (2016); \n"
                              "* Exp2: the improved estimator introduced in Cordero-Grande et al. (2019); \n"
-                             "* Import: import from a pre-estimated noise level map; \n"
                              "* Med: estimate based on the median eigenvalue as in Gavish and Donohue (2014); \n"
                              "* MRM2022: the alternative estimator introduced in Olesen et al. (2022).") +
                       Argument("algorithm").type_choice(estimators);
 
-std::shared_ptr<Base> make_estimator() {
-  auto opt = App::get_options("estimator");
-  const estimator_type est = opt.empty() ? estimator_type::EXP2 : estimator_type((int)(opt[0][0]));
-  auto noise_in = get_options("noise_in");
-  // TODO Remove once -noise_in is used in other ways
-  if (!noise_in.empty() && est != estimator_type::IMPORT) {
-    WARN("-noise_in option has no effect unless -estimator import is specified");
+std::shared_ptr<Base> make_estimator(const bool permit_noise_in) {
+  auto opt = get_options("estimator");
+  if (permit_noise_in) {
+    auto noise_in = get_options("noise_in");
+    if (!noise_in.empty()) {
+      if (!opt.empty())
+        throw Exception("Cannot both provide an input noise level image and specify a noise level estimator");
+      return std::make_shared<Import>(noise_in[0][0]);
+    }
   }
+  const estimator_type est = opt.empty() ? estimator_type::EXP2 : estimator_type((int)(opt[0][0]));
   switch (est) {
   case estimator_type::EXP1:
     return std::make_shared<Exp<1>>();
   case estimator_type::EXP2:
     return std::make_shared<Exp<2>>();
-  case estimator_type::IMPORT:
-    if (noise_in.empty())
-      throw Exception("-estimator import requires a pre-estimated noise level image "
-                      "to be provided via the -noise_in option");
-    return std::make_shared<Import>(noise_in[0][0]);
   case estimator_type::MED:
     return std::make_shared<Med>();
   case estimator_type::MRM2022:
diff --git a/src/denoise/estimator/estimator.h b/src/denoise/estimator/estimator.h
index 10cae55f43..0a1cac5c8e 100644
--- a/src/denoise/estimator/estimator.h
+++ b/src/denoise/estimator/estimator.h
@@ -27,8 +27,8 @@ namespace MR::Denoise::Estimator {
 class Base;
 
 extern const App::Option option;
-const std::vector<std::string> estimators = {"exp1", "exp2", "import", "med", "mrm2022"};
-enum class estimator_type { EXP1, EXP2, IMPORT, MED, MRM2022 };
-std::shared_ptr<Base> make_estimator();
+const std::vector<std::string> estimators = {"exp1", "exp2", "med", "mrm2022"};
+enum class estimator_type { EXP1, EXP2, MED, MRM2022 };
+std::shared_ptr<Base> make_estimator(const bool permit_noise_in);
 
 } // namespace MR::Denoise::Estimator
diff --git a/src/denoise/kernel/base.h b/src/denoise/kernel/base.h
index 0e8d47b7ed..490bad437c 100644
--- a/src/denoise/kernel/base.h
+++ b/src/denoise/kernel/base.h
@@ -19,12 +19,18 @@
 #include "denoise/kernel/data.h"
 #include "denoise/kernel/voxel.h"
 #include "header.h"
+#include "transform.h"
 
 namespace MR::Denoise::Kernel {
 
 class Base {
 public:
-  Base(const Header &H) : H(H) {}
+  Base(const Header &H, const std::array<ssize_t, 3> &subsample_factors)
+      : H(H),
+        transform(H),
+        halfvoxel_offsets({subsample_factors[0] & 1 ? 0.0 : 0.5,
+                           subsample_factors[1] & 1 ? 0.0 : 0.5,
+                           subsample_factors[2] & 1 ? 0.0 : 0.5}) {}
   Base(const Base &) = default;
   virtual ~Base() = default;
   // This is just for pre-allocating matrices
@@ -34,6 +40,20 @@ class Base {
 
 protected:
   const Header H;
+  const Transform transform;
+  std::array<default_type, 3> halfvoxel_offsets;
+
+  // For translating the index of a processed voxel
+  //   into a realspace position corresponding to the centre of the patch,
+  //   accounting for the fact that subsampling may be introducing an offset
+  //   such that the actual centre of the patch is not at the centre of this voxel
+  Eigen::Vector3d voxel2real(const Kernel::Voxel::index_type &pos) const {
+    return (                                               //
+        transform.voxel2scanner *                          //
+        Eigen::Vector3d({pos[0] + halfvoxel_offsets[0],    //
+                         pos[1] + halfvoxel_offsets[1],    //
+                         pos[2] + halfvoxel_offsets[2]})); //
+  }
 };
 
 } // namespace MR::Denoise::Kernel
diff --git a/src/denoise/kernel/cuboid.cpp b/src/denoise/kernel/cuboid.cpp
index b65b579a3a..06aefa30d8 100644
--- a/src/denoise/kernel/cuboid.cpp
+++ b/src/denoise/kernel/cuboid.cpp
@@ -19,9 +19,9 @@
 namespace MR::Denoise::Kernel {
 
 Cuboid::Cuboid(const Header &header,
-               const std::array<ssize_t, 3> &extent,
-               const std::array<ssize_t, 3> &subsample_factors)
-    : Base(header),
+               const std::array<ssize_t, 3> &subsample_factors,
+               const std::array<ssize_t, 3> &extent)
+    : Base(header, subsample_factors),
       size(extent[0] * extent[1] * extent[2]),
       // Only sensible if no subsampling is performed,
       //   and every single DWI voxel is reconstructed from a patch centred at that voxel,
@@ -34,14 +34,12 @@ Cuboid::Cuboid(const Header &header,
                         "size of cubic kernel must be an odd integer");
       bounding_box(axis, 0) = -extent[axis] / 2;
       bounding_box(axis, 1) = extent[axis] / 2;
-      halfvoxel_offsets[axis] = 0.0;
     } else {
       if (extent[axis] % 2)
         throw Exception("For subsampling by an even number, "
                         "size of cubic kernel must be an even integer");
       bounding_box(axis, 0) = 1 - extent[axis] / 2;
       bounding_box(axis, 1) = extent[axis] / 2;
-      halfvoxel_offsets[axis] = 0.5;
     }
   }
 }
@@ -59,7 +57,7 @@ inline ssize_t wrapindex(int p, int r, int bbminus, int bbplus, int max) {
 } // namespace
 
 Data Cuboid::operator()(const Voxel::index_type &pos) const {
-  Data result(centre_index);
+  Data result(voxel2real(pos), centre_index);
   Voxel::index_type voxel;
   Offset::index_type offset;
   for (offset[2] = bounding_box(2, 0); offset[2] <= bounding_box(2, 1); ++offset[2]) {
diff --git a/src/denoise/kernel/cuboid.h b/src/denoise/kernel/cuboid.h
index 93a318e136..c90d915efd 100644
--- a/src/denoise/kernel/cuboid.h
+++ b/src/denoise/kernel/cuboid.h
@@ -27,7 +27,7 @@ namespace MR::Denoise::Kernel {
 class Cuboid : public Base {
 
 public:
-  Cuboid(const Header &header, const std::array<ssize_t, 3> &extent, const std::array<ssize_t, 3> &subsample_factors);
+  Cuboid(const Header &header, const std::array<ssize_t, 3> &subsample_factors, const std::array<ssize_t, 3> &extent);
   Cuboid(const Cuboid &) = default;
   ~Cuboid() final = default;
   Data operator()(const Voxel::index_type &pos) const override;
@@ -37,7 +37,6 @@ class Cuboid : public Base {
   Eigen::Array<int, 3, 2> bounding_box;
   const ssize_t size;
   const ssize_t centre_index;
-  std::array<default_type, 3> halfvoxel_offsets;
 };
 
 } // namespace MR::Denoise::Kernel
diff --git a/src/denoise/kernel/data.h b/src/denoise/kernel/data.h
index 5d783baa05..f4659d21a6 100644
--- a/src/denoise/kernel/data.h
+++ b/src/denoise/kernel/data.h
@@ -18,6 +18,8 @@
 
 #include <vector>
 
+#include <Eigen/Dense>
+
 #include "denoise/kernel/voxel.h"
 #include "types.h"
 
@@ -25,11 +27,26 @@ namespace MR::Denoise::Kernel {
 
 class Data {
 public:
-  Data() : centre_index(-1), max_distance(-std::numeric_limits<default_type>::infinity()) {}
-  Data(const ssize_t i) : centre_index(i), max_distance(-std::numeric_limits<default_type>::infinity()) {}
+  Data(const Eigen::Vector3d &pos, const ssize_t i)
+      : centre_realspace(pos),
+        centre_index(i),
+        max_distance(-std::numeric_limits<default_type>::infinity()),
+        centre_noise(std::numeric_limits<double>::signaling_NaN()) {}
+  Data(const ssize_t i)
+      : centre_realspace(Eigen::Vector3d::Constant(std::numeric_limits<double>::signaling_NaN())),
+        centre_index(i),
+        max_distance(-std::numeric_limits<default_type>::infinity()),
+        centre_noise(std::numeric_limits<double>::signaling_NaN()) {}
+  Data()
+      : centre_realspace(Eigen::Vector3d::Constant(std::numeric_limits<double>::signaling_NaN())),
+        centre_index(-1),
+        max_distance(-std::numeric_limits<default_type>::infinity()),
+        centre_noise(std::numeric_limits<double>::signaling_NaN()) {}
+  Eigen::Vector3d centre_realspace;
   std::vector<Voxel> voxels;
   ssize_t centre_index;
   default_type max_distance;
+  double centre_noise;
 };
 
 } // namespace MR::Denoise::Kernel
diff --git a/src/denoise/kernel/kernel.cpp b/src/denoise/kernel/kernel.cpp
index e991b86508..f935bf13b6 100644
--- a/src/denoise/kernel/kernel.cpp
+++ b/src/denoise/kernel/kernel.cpp
@@ -96,8 +96,8 @@ std::shared_ptr<Base> make_kernel(const Header &H, const std::array<ssize_t, 3>
     opt = get_options("radius_mm");
     if (opt.empty())
       return std::make_shared<SphereRatio>(
-          H, get_option_value("radius_ratio", sphere_multiplier_default), subsample_factors);
-    return std::make_shared<SphereFixedRadius>(H, opt[0][0], subsample_factors);
+          H, subsample_factors, get_option_value("radius_ratio", sphere_multiplier_default));
+    return std::make_shared<SphereFixedRadius>(H, subsample_factors, opt[0][0]);
   }
   case Kernel::shape_type::CUBOID: {
     if (!get_options("radius_mm").empty() || !get_options("radius_ratio").empty())
@@ -142,7 +142,7 @@ std::shared_ptr<Base> make_kernel(const Header &H, const std::array<ssize_t, 3>
            "and cause inconsistent denoising between adjacent voxels");
     }
 
-    return std::make_shared<Cuboid>(H, extent, subsample_factors);
+    return std::make_shared<Cuboid>(H, subsample_factors, extent);
   } break;
   default:
     assert(false);
diff --git a/src/denoise/kernel/sphere_base.cpp b/src/denoise/kernel/sphere_base.cpp
index 666b090220..11216a994d 100644
--- a/src/denoise/kernel/sphere_base.cpp
+++ b/src/denoise/kernel/sphere_base.cpp
@@ -21,20 +21,18 @@
 namespace MR::Denoise::Kernel {
 
 SphereBase::Shared::Shared(const Header &voxel_grid,
-                           const default_type max_radius,
-                           const std::array<ssize_t, 3> &subsample_factors) {
+                           const std::array<ssize_t, 3> &subsample_factors,
+                           const std::array<default_type, 3> &halfvoxel_offsets,
+                           const default_type max_radius) {
   const default_type max_radius_sq = Math::pow2(max_radius);
   Eigen::Array<int, 3, 2> bounding_box;
-  std::array<default_type, 3> halfvoxel_offsets;
   for (ssize_t axis = 0; axis != 3; ++axis) {
     if (subsample_factors[axis] % 2) {
       bounding_box(axis, 1) = int(std::ceil(max_radius / voxel_grid.spacing(axis)));
       bounding_box(axis, 0) = -bounding_box(axis, 1);
-      halfvoxel_offsets[axis] = 0.0;
     } else {
       bounding_box(axis, 0) = -int(std::ceil((max_radius / voxel_grid.spacing(axis)) - 0.5));
       bounding_box(axis, 1) = int(std::ceil((max_radius / voxel_grid.spacing(axis)) + 0.5));
-      halfvoxel_offsets[axis] = 0.5;
     }
   }
   // Build the searchlight
diff --git a/src/denoise/kernel/sphere_base.h b/src/denoise/kernel/sphere_base.h
index da1bfa4871..6d8dc5ace1 100644
--- a/src/denoise/kernel/sphere_base.h
+++ b/src/denoise/kernel/sphere_base.h
@@ -30,8 +30,10 @@ namespace MR::Denoise::Kernel {
 class SphereBase : public Base {
 
 public:
-  SphereBase(const Header &voxel_grid, const default_type max_radius, const std::array<ssize_t, 3> &subsample_factors)
-      : Base(voxel_grid), shared(new Shared(voxel_grid, max_radius, subsample_factors)) {}
+  SphereBase(const Header &voxel_grid, const std::array<ssize_t, 3> &subsample_factors, const default_type max_radius)
+      : Base(voxel_grid, subsample_factors),
+        shared(new Shared(voxel_grid, subsample_factors, halfvoxel_offsets, max_radius)),
+        centre_index(subsample_factors == std::array<ssize_t, 3>({1, 1, 1}) ? 0 : -1) {}
 
   SphereBase(const SphereBase &) = default;
 
@@ -41,7 +43,10 @@ class SphereBase : public Base {
   class Shared {
   public:
     using TableType = std::vector<Offset>;
-    Shared(const Header &voxel_grid, const default_type max_radius, const std::array<ssize_t, 3> &subsample_factors);
+    Shared(const Header &voxel_grid,
+           const std::array<ssize_t, 3> &subsample_factors,
+           const std::array<default_type, 3> &halfvoxel_offsets,
+           const default_type max_radius);
     TableType::const_iterator begin() const { return data.begin(); }
     TableType::const_iterator end() const { return data.end(); }
 
@@ -50,6 +55,7 @@ class SphereBase : public Base {
   };
 
   std::shared_ptr<Shared> shared;
+  const ssize_t centre_index;
 };
 
 } // namespace MR::Denoise::Kernel
diff --git a/src/denoise/kernel/sphere_radius.cpp b/src/denoise/kernel/sphere_radius.cpp
index d759ce715b..be5e4cc1e5 100644
--- a/src/denoise/kernel/sphere_radius.cpp
+++ b/src/denoise/kernel/sphere_radius.cpp
@@ -19,7 +19,7 @@
 namespace MR::Denoise::Kernel {
 
 Data SphereFixedRadius::operator()(const Voxel::index_type &pos) const {
-  Data result(0);
+  Data result(voxel2real(pos), centre_index);
   result.voxels.reserve(maximum_size);
   for (auto map_it = shared->begin(); map_it != shared->end(); ++map_it) {
     const Voxel::index_type voxel({pos[0] + map_it->index[0],   //
diff --git a/src/denoise/kernel/sphere_radius.h b/src/denoise/kernel/sphere_radius.h
index 73a533d022..1ebb0607c9 100644
--- a/src/denoise/kernel/sphere_radius.h
+++ b/src/denoise/kernel/sphere_radius.h
@@ -26,9 +26,9 @@ namespace MR::Denoise::Kernel {
 class SphereFixedRadius : public SphereBase {
 public:
   SphereFixedRadius(const Header &voxel_grid,                         //
-                    const default_type radius,                        //
-                    const std::array<ssize_t, 3> &subsample_factors)  //
-      : SphereBase(voxel_grid, radius, subsample_factors),            //
+                    const std::array<ssize_t, 3> &subsample_factors,  //
+                    const default_type radius)                        //
+      : SphereBase(voxel_grid, subsample_factors, radius),            //
         maximum_size(std::distance(shared->begin(), shared->end())) { //
     INFO("Maximum number of voxels in " + str(radius) + "mm fixed-radius kernel is " + str(maximum_size));
   }
diff --git a/src/denoise/kernel/sphere_ratio.cpp b/src/denoise/kernel/sphere_ratio.cpp
index f682c37198..f4532fc002 100644
--- a/src/denoise/kernel/sphere_ratio.cpp
+++ b/src/denoise/kernel/sphere_ratio.cpp
@@ -19,7 +19,7 @@
 namespace MR::Denoise::Kernel {
 
 Data SphereRatio::operator()(const Voxel::index_type &pos) const {
-  Data result(0);
+  Data result(voxel2real(pos), centre_index);
   auto table_it = shared->begin();
   while (table_it != shared->end()) {
     // If there's a tie in distances, want to include all such offsets in the kernel,
diff --git a/src/denoise/kernel/sphere_ratio.h b/src/denoise/kernel/sphere_ratio.h
index b6334bc884..36778ac2e4 100644
--- a/src/denoise/kernel/sphere_ratio.h
+++ b/src/denoise/kernel/sphere_ratio.h
@@ -27,8 +27,8 @@ constexpr default_type sphere_multiplier_default = 1.0 / 0.85;
 class SphereRatio : public SphereBase {
 
 public:
-  SphereRatio(const Header &voxel_grid, const default_type min_ratio, const std::array<ssize_t, 3> &subsample_factors)
-      : SphereBase(voxel_grid, compute_max_radius(voxel_grid, min_ratio), subsample_factors),
+  SphereRatio(const Header &voxel_grid, const std::array<ssize_t, 3> &subsample_factors, const default_type min_ratio)
+      : SphereBase(voxel_grid, subsample_factors, compute_max_radius(voxel_grid, min_ratio)),
         min_size(std::ceil(voxel_grid.size(3) * min_ratio)) {}
 
   SphereRatio(const SphereRatio &) = default;
diff --git a/src/denoise/kernel/voxel.h b/src/denoise/kernel/voxel.h
index c76ff286ce..0570d2c8d7 100644
--- a/src/denoise/kernel/voxel.h
+++ b/src/denoise/kernel/voxel.h
@@ -48,7 +48,20 @@ template <class T> class VoxelBase {
 
 // Need signed integer to represent offsets from the centre of the kernel;
 //   however absolute voxel indices should be unsigned
-using Voxel = VoxelBase<ssize_t>;
+// For nonstationarity correction,
+//   "Voxel" also needs a noise level estimate per voxel in the patch,
+//   as the scaling on insertion of data into the matrix
+//   needs to be reversed when reconstructing the denoised signal from the eigenvectors
+class Voxel : public VoxelBase<ssize_t> {
+public:
+  using index_type = VoxelBase<ssize_t>::index_type;
+  Voxel(const index_type &index, const default_type sq_distance, const default_type noise_level)
+      : VoxelBase<ssize_t>(index, sq_distance), noise_level(noise_level) {}
+  Voxel(const index_type &index, const default_type sq_distance)
+      : VoxelBase<ssize_t>(index, sq_distance), noise_level(std::numeric_limits<default_type>::signaling_NaN()) {}
+  default_type noise_level;
+};
+
 using Offset = VoxelBase<int>;
 
 } // namespace MR::Denoise::Kernel
diff --git a/src/denoise/recon.cpp b/src/denoise/recon.cpp
index 4eb5df51e3..b428d94f3f 100644
--- a/src/denoise/recon.cpp
+++ b/src/denoise/recon.cpp
@@ -24,11 +24,12 @@ template <typename F>
 Recon<F>::Recon(const Header &header,
                 std::shared_ptr<Subsample> subsample,
                 std::shared_ptr<Kernel::Base> kernel,
+                Image<float> &nonstationarity_image,
                 std::shared_ptr<Estimator::Base> estimator,
                 filter_type filter,
                 aggregator_type aggregator,
                 Exports &exports)
-    : Estimate<F>(header, subsample, kernel, estimator, exports),
+    : Estimate<F>(header, subsample, kernel, nonstationarity_image, estimator, exports),
       filter(filter),
       aggregator(aggregator),
       // FWHM = 2 x cube root of spacings between kernels
@@ -46,7 +47,7 @@ template <typename F> void Recon<F>::operator()(Image<F> &dwi, Image<F> &out) {
 
   Estimate<F>::operator()(dwi);
 
-  const ssize_t n = Estimate<F>::neighbourhood.voxels.size();
+  const ssize_t n = Estimate<F>::patch.voxels.size();
   const ssize_t r = std::min(Estimate<F>::m, n);
   const ssize_t q = std::max(Estimate<F>::m, n);
   const double beta = double(r) / double(q);
@@ -136,18 +137,19 @@ template <typename F> void Recon<F>::operator()(Image<F> &dwi, Image<F> &out) {
   //   then we need to instead compute the full projection
   switch (aggregator) {
   case aggregator_type::EXCLUSIVE:
+    assert(Estimate<F>::patch.centre_index >= 0);
     if (Estimate<F>::m <= n)
-      Xr.noalias() =                                                       //
-          Estimate<F>::eig.eigenvectors() *                                //
-          (w.head(r).cast<F>().matrix().asDiagonal() *                     //
-           (Estimate<F>::eig.eigenvectors().adjoint() *                    //
-            Estimate<F>::X.col(Estimate<F>::neighbourhood.centre_index))); //
+      Xr.noalias() =                                               //
+          Estimate<F>::eig.eigenvectors() *                        //
+          (w.head(r).cast<F>().matrix().asDiagonal() *             //
+           (Estimate<F>::eig.eigenvectors().adjoint() *            //
+            Estimate<F>::X.col(Estimate<F>::patch.centre_index))); //
     else
-      Xr.noalias() =                                                                                  //
-          Estimate<F>::X.leftCols(n) *                                                                //
-          (Estimate<F>::eig.eigenvectors() *                                                          //
-           (w.head(r).cast<F>().matrix().asDiagonal() *                                               //
-            Estimate<F>::eig.eigenvectors().adjoint().col(Estimate<F>::neighbourhood.centre_index))); //
+      Xr.noalias() =                                                                          //
+          Estimate<F>::X.leftCols(n) *                                                        //
+          (Estimate<F>::eig.eigenvectors() *                                                  //
+           (w.head(r).cast<F>().matrix().asDiagonal() *                                       //
+            Estimate<F>::eig.eigenvectors().adjoint().col(Estimate<F>::patch.centre_index))); //
     assert(Xr.allFinite());
     assign_pos_of(dwi).to(out);
     out.row(3) = Xr.col(0);
@@ -177,17 +179,22 @@ template <typename F> void Recon<F>::operator()(Image<F> &dwi, Image<F> &out) {
             Estimate<F>::eig.eigenvectors().adjoint())); //
     }
     assert(Xr.leftCols(n).allFinite());
+    // Undo prior within-patch non-stationarity correction
+    if (std::isfinite(Estimate<F>::patch.centre_noise)) {
+      for (ssize_t i = 0; i != n; ++i)
+        Xr.col(i) *= Estimate<F>::patch.voxels[i].noise_level / Estimate<F>::patch.centre_noise;
+    }
     std::lock_guard<std::mutex> lock(Estimate<F>::mutex);
-    for (size_t voxel_index = 0; voxel_index != Estimate<F>::neighbourhood.voxels.size(); ++voxel_index) {
-      assign_pos_of(Estimate<F>::neighbourhood.voxels[voxel_index].index, 0, 3).to(out);
-      assign_pos_of(Estimate<F>::neighbourhood.voxels[voxel_index].index).to(Estimate<F>::exports.sum_aggregation);
+    for (size_t voxel_index = 0; voxel_index != Estimate<F>::patch.voxels.size(); ++voxel_index) {
+      assign_pos_of(Estimate<F>::patch.voxels[voxel_index].index, 0, 3).to(out);
+      assign_pos_of(Estimate<F>::patch.voxels[voxel_index].index).to(Estimate<F>::exports.sum_aggregation);
       double weight = std::numeric_limits<double>::signaling_NaN();
       switch (aggregator) {
       case aggregator_type::EXCLUSIVE:
         assert(false);
         break;
       case aggregator_type::GAUSSIAN:
-        weight = std::exp(gaussian_multiplier * Estimate<F>::neighbourhood.voxels[voxel_index].sq_distance);
+        weight = std::exp(gaussian_multiplier * Estimate<F>::patch.voxels[voxel_index].sq_distance);
         break;
       case aggregator_type::INVL0:
         weight = 1.0 / (1 + out_rank);
@@ -202,7 +209,7 @@ template <typename F> void Recon<F>::operator()(Image<F> &dwi, Image<F> &out) {
       out.row(3) += weight * Xr.col(voxel_index);
       Estimate<F>::exports.sum_aggregation.value() += weight;
       if (Estimate<F>::exports.rank_output.valid()) {
-        assign_pos_of(Estimate<F>::neighbourhood.voxels[voxel_index].index, 0, 3).to(Estimate<F>::exports.rank_output);
+        assign_pos_of(Estimate<F>::patch.voxels[voxel_index].index, 0, 3).to(Estimate<F>::exports.rank_output);
         Estimate<F>::exports.rank_output.value() += weight * out_rank;
       }
     }
diff --git a/src/denoise/recon.h b/src/denoise/recon.h
index 6db9e40d59..977df221f4 100644
--- a/src/denoise/recon.h
+++ b/src/denoise/recon.h
@@ -37,6 +37,7 @@ template <typename F> class Recon : public Estimate<F> {
   Recon(const Header &header,
         std::shared_ptr<Subsample> subsample,
         std::shared_ptr<Kernel::Base> kernel,
+        Image<float> &nonstationarity_image,
         std::shared_ptr<Estimator::Base> estimator,
         filter_type filter,
         aggregator_type aggregator,
diff --git a/src/denoise/subsample.cpp b/src/denoise/subsample.cpp
index 32d9280fbc..92f5c6b8b7 100644
--- a/src/denoise/subsample.cpp
+++ b/src/denoise/subsample.cpp
@@ -105,12 +105,4 @@ Header Subsample::make_subsample_header() const {
   return H;
 }
 
-Eigen::Vector3d Subsample::patch_centre(const Kernel::Voxel::index_type &pos) const {
-  assert(process(pos));
-  const Eigen::Vector3d centre({double(pos[0]) + (factors[0] & 1 ? 0.0 : 0.5),
-                                double(pos[1]) + (factors[1] & 1 ? 0.0 : 0.5),
-                                double(pos[2]) + (factors[2] & 1 ? 0.0 : 0.5)});
-  return H_in.transform() * centre;
-}
-
 } // namespace MR::Denoise
diff --git a/src/denoise/subsample.h b/src/denoise/subsample.h
index d5b8700cd9..825de3b88b 100644
--- a/src/denoise/subsample.h
+++ b/src/denoise/subsample.h
@@ -33,10 +33,6 @@ class Subsample {
 
   const Header &header() const { return H_ss; }
 
-  // TODO What other functionalities does this class need?
-  // - Decide whether a given 3-vector voxel position should be sampled vs. not
-  // - Convert input image 3-vector voxel position to output subsampled header voxel position
-  // - Convert subsampled header voxel position to centre voxel in input image
   // TODO May want to move definition of Kernel::Voxel out of Kernel namespace
   bool process(const Kernel::Voxel::index_type &pos) const;
   std::array<ssize_t, 3> in2ss(const Kernel::Voxel::index_type &pos) const;
@@ -45,10 +41,6 @@ class Subsample {
 
   static std::shared_ptr<Subsample> make(const Header &in);
 
-  // TODO From a processed input voxel,
-  //   get the realspace position of the centre of the patch
-  Eigen::Vector3d patch_centre(const Kernel::Voxel::index_type &pos) const;
-
 protected:
   const Header H_in;
   const std::array<ssize_t, 3> factors;

From 83ebb1f965d4f67f5067a12dd13ecffd8cfc7543 Mon Sep 17 00:00:00 2001
From: Robert Smith <robert.smith@florey.edu.au>
Date: Wed, 4 Dec 2024 13:59:36 +1100
Subject: [PATCH 27/34] dwidenoise: Add missing file for pre-estimated noise
 map

File missing from bc33d4d61348d9357d72114896b911b683cf0602.
---
 src/denoise/estimator/import.h | 76 ++++++++++++++++++++++++++++++++++
 1 file changed, 76 insertions(+)
 create mode 100644 src/denoise/estimator/import.h

diff --git a/src/denoise/estimator/import.h b/src/denoise/estimator/import.h
new file mode 100644
index 0000000000..546a6c1cd1
--- /dev/null
+++ b/src/denoise/estimator/import.h
@@ -0,0 +1,76 @@
+/* Copyright (c) 2008-2024 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#pragma once
+
+#include <string>
+
+#include "denoise/denoise.h"
+#include "denoise/estimator/base.h"
+#include "denoise/estimator/result.h"
+#include "image.h"
+#include "interp/cubic.h"
+
+namespace MR::Denoise::Estimator {
+
+class Import : public Base {
+public:
+  Import(const std::string &path) : noise_image(Image<float>::open(path)) {}
+  Result operator()(const eigenvalues_type &s,                //
+                    const ssize_t m,                          //
+                    const ssize_t n,                          //
+                    const Eigen::Vector3d &pos) const final { //
+    Result result;
+    const ssize_t r = std::min(m, n);
+    const ssize_t q = std::max(m, n);
+    {
+      // Construct on each call to preserve const-ness & thread-safety
+      Interp::Cubic<Image<float>> interp(noise_image);
+      // TODO This will cause issues at the edge of the image FoV
+      // Addressing this may require integration of the mrfilter changes
+      //   that provide wrappers for various handling of FoV edges
+      // For now, just expect that denoising won't do anything
+      //   where the patch centre is too close to the image edge for cubic interpolation
+      if (!interp.scanner(pos))
+        return result;
+      result.sigma2 = Math::pow2(interp.value());
+    }
+    // From this noise level,
+    //   estimate the upper bound of the MP distribution and rank of signal
+    //   given the ordered list of eigenvalues
+    double cumulative_lambda = 0.0;
+    double recalc_sigmasq = 0.0;
+    for (ssize_t p = 0; p != r; ++p) {
+      const double lambda = std::max(s[p], 0.0) / q;
+      cumulative_lambda += lambda;
+      const double sigma_sq = cumulative_lambda / (p + 1);
+      if (sigma_sq < result.sigma2) {
+        result.cutoff_p = p;
+        result.lamplus = lambda;
+        recalc_sigmasq = sigma_sq;
+      }
+    }
+    // TODO It would be nice if the upper bound, lambda_plus,
+    //   could be yielded at a higher precision than the discrete eigenvalues,
+    //   as optimal shrinkage / optimal thresholding could make use of this precision if available
+    return result;
+  }
+
+private:
+  Image<float> noise_image;
+};
+
+} // namespace MR::Denoise::Estimator

From bba4fe17b4eaa152adc2cb152161f5bd73bf740c Mon Sep 17 00:00:00 2001
From: Robert Smith <robert.smith@florey.edu.au>
Date: Wed, 4 Dec 2024 14:28:38 +1100
Subject: [PATCH 28/34] Non-linear phase demodulation

- dwi2noise and dwidenoise: Ability to estimate and remove non-linear phase across images prior to denoising. In the case of dwidenoise this phase is re-introduced into the data after PCA denoising. This is now the default over linear phase ramp removal.
- mrfilter: New filter "kspace"; this currently only supports a Tukey window, therefore demonstrating the operation of the demodulation that is in use for DWI denoising, but can be used as a standalone operation, and is intended to be extensible to incorporate other k-space filter designs.
---
 cmd/dwi2noise.cpp          |  22 ++--
 cmd/dwidenoise.cpp         |  15 ++-
 cmd/mrfilter.cpp           |  65 ++++++++++--
 core/filter/demodulate.h   |  55 ++++++++--
 core/filter/kspace.h       | 200 +++++++++++++++++++++++++++++++++++++
 src/denoise/demodulate.cpp | 117 +++++++++++++---------
 src/denoise/demodulate.h   |  16 ++-
 7 files changed, 403 insertions(+), 87 deletions(-)
 create mode 100644 core/filter/kspace.h

diff --git a/cmd/dwi2noise.cpp b/cmd/dwi2noise.cpp
index f64bc25333..6c2eb17e62 100644
--- a/cmd/dwi2noise.cpp
+++ b/cmd/dwi2noise.cpp
@@ -169,20 +169,20 @@ void run(Header &data,
 
 template <typename T>
 void run(Header &data,
-         const std::vector<size_t> &demodulation_axes,
+         const Demodulation &demodulation,
          std::shared_ptr<Subsample> subsample,
          std::shared_ptr<Kernel::Base> kernel,
          Image<float> &nonstationarity_image,
          std::shared_ptr<Estimator::Base> estimator,
          Exports &exports) {
-  if (demodulation_axes.empty()) {
+  if (!demodulation) {
     run<T>(data, subsample, kernel, nonstationarity_image, estimator, exports);
     return;
   }
   auto input = data.get_image<T>();
   auto input_demod = Image<T>::scratch(data, "Phase-demodulated version of \"" + data.name() + "\"");
   {
-    Filter::Demodulate demodulator(input, demodulation_axes);
+    Filter::Demodulate demodulator(input, demodulation.axes, demodulation.mode == demodulation_t::LINEAR);
     demodulator(input, input_demod);
   }
   Estimate<T> func(data, subsample, kernel, nonstationarity_image, estimator, exports);
@@ -195,10 +195,7 @@ void run() {
     throw Exception("input image must be 4-dimensional");
   bool complex = dwi.datatype().is_complex();
 
-  Image<float> nonstationarity_image;
-  auto opt = get_options("nonstationarity");
-  if (!opt.empty())
-    nonstationarity_image = Image<float>::open(opt[0][0]);
+  const Demodulation demodulation = get_demodulation(dwi);
 
   auto subsample = Subsample::make(dwi);
   assert(subsample);
@@ -206,6 +203,11 @@ void run() {
   auto kernel = Kernel::make_kernel(dwi, subsample->get_factors());
   assert(kernel);
 
+  Image<float> nonstationarity_image;
+  auto opt = get_options("nonstationarity");
+  if (!opt.empty())
+    nonstationarity_image = Image<float>::open(opt[0][0]);
+
   auto estimator = Estimator::make_estimator(false);
   assert(estimator);
 
@@ -224,8 +226,6 @@ void run() {
   if (!opt.empty())
     exports.set_patchcount(opt[0][0]);
 
-  const std::vector<size_t> demodulation_axes = get_demodulation_axes(dwi);
-
   int prec = get_option_value("datatype", 0); // default: single precision
   if (complex)
     prec += 2; // support complex input data
@@ -242,11 +242,11 @@ void run() {
     break;
   case 2:
     INFO("select complex float32 for processing");
-    run<cfloat>(dwi, demodulation_axes, subsample, kernel, nonstationarity_image, estimator, exports);
+    run<cfloat>(dwi, demodulation, subsample, kernel, nonstationarity_image, estimator, exports);
     break;
   case 3:
     INFO("select complex float64 for processing");
-    run<cdouble>(dwi, demodulation_axes, subsample, kernel, nonstationarity_image, estimator, exports);
+    run<cdouble>(dwi, demodulation, subsample, kernel, nonstationarity_image, estimator, exports);
     break;
   }
 }
diff --git a/cmd/dwidenoise.cpp b/cmd/dwidenoise.cpp
index c50c90a452..de9c31dd79 100644
--- a/cmd/dwidenoise.cpp
+++ b/cmd/dwidenoise.cpp
@@ -273,7 +273,7 @@ void run(Header &data,
 
 template <typename T>
 void run(Header &data,
-         const std::vector<size_t> &demodulation_axes,
+         const Demodulation &demodulation,
          std::shared_ptr<Subsample> subsample,
          std::shared_ptr<Kernel::Base> kernel,
          Image<float> &nonstationarity_image,
@@ -282,7 +282,7 @@ void run(Header &data,
          aggregator_type aggregator,
          const std::string &output_name,
          Exports &exports) {
-  if (demodulation_axes.empty()) {
+  if (!demodulation) {
     run<T>(data, subsample, kernel, nonstationarity_image, estimator, filter, aggregator, output_name, exports);
     return;
   }
@@ -293,7 +293,7 @@ void run(Header &data,
   H_scratch.datatype() = DataType::from<T>();
   H_scratch.datatype().set_byte_order_native();
   auto input_demodulated = Image<T>::scratch(H_scratch, "Phase-demodulated version of input DWI");
-  Filter::Demodulate demodulate(input, demodulation_axes);
+  Filter::Demodulate demodulate(input, demodulation.axes, demodulation.mode == demodulation_t::LINEAR);
   demodulate(input, input_demodulated, false);
   input = Image<T>(); // free memory
   // create output
@@ -320,10 +320,11 @@ void run(Header &data,
 
 void run() {
   auto dwi = Header::open(argument[0]);
-
   if (dwi.ndim() != 4 || dwi.size(3) <= 1)
     throw Exception("input image must be 4-dimensional");
 
+  const Demodulation demodulation = get_demodulation(dwi);
+
   auto subsample = Subsample::make(dwi);
   assert(subsample);
 
@@ -398,8 +399,6 @@ void run() {
     exports.set_sum_aggregation("");
   }
 
-  const std::vector<size_t> demodulation_axes = get_demodulation_axes(dwi);
-
   int prec = get_option_value("datatype", 0); // default: single precision
   if (dwi.datatype().is_complex())
     prec += 2; // support complex input data
@@ -417,7 +416,7 @@ void run() {
   case 2:
     INFO("select complex float32 for processing");
     run<cfloat>(dwi,
-                demodulation_axes,
+                demodulation,
                 subsample,
                 kernel,
                 nonstationarity_image,
@@ -430,7 +429,7 @@ void run() {
   case 3:
     INFO("select complex float64 for processing");
     run<cdouble>(dwi,
-                 demodulation_axes,
+                 demodulation,
                  subsample,
                  kernel,
                  nonstationarity_image,
diff --git a/cmd/mrfilter.cpp b/cmd/mrfilter.cpp
index 4e37def026..78d06c00ab 100644
--- a/cmd/mrfilter.cpp
+++ b/cmd/mrfilter.cpp
@@ -21,6 +21,7 @@
 #include "filter/base.h"
 #include "filter/demodulate.h"
 #include "filter/gradient.h"
+#include "filter/kspace.h"
 #include "filter/median.h"
 #include "filter/normalise.h"
 #include "filter/smooth.h"
@@ -32,15 +33,20 @@
 using namespace MR;
 using namespace App;
 
-const std::vector<std::string> filters = {"demodulate", "fft", "gradient", "median", "smooth", "normalise", "zclean"};
+const std::vector<std::string> filters = {
+    "demodulate", "fft", "gradient", "kspace", "median", "smooth", "normalise", "zclean"};
 
 // clang-format off
-const OptionGroup FFTAxesOption = OptionGroup ("Options applicable to both demodulate and FFT filters")
+const OptionGroup FFTAxesOption = OptionGroup ("Options applicable to demodulate / FFT / k-space filters")
   + Option ("axes", "the axes along which to apply the Fourier Transform."
                     " By default, the transform is applied along the three spatial axes."
                     " Provide as a comma-separate list of axis indices.")
     + Argument ("list").type_sequence_int();
 
+const OptionGroup DemodulateOption = OptionGroup ("Options applicable to demodulate filter")
+  + Option ("linear", "only demodulate based on a linear phase ramp, "
+                      "rather than a filtered k-space");
+
 const OptionGroup FFTOption = OptionGroup ("Options for FFT filter")
   + Option ("inverse", "apply the inverse FFT")
   + Option ("magnitude", "output a magnitude image rather than a complex-valued image")
@@ -62,6 +68,15 @@ const OptionGroup GradientOption = OptionGroup ("Options for gradient filter")
   + Option ("scanner", "define the gradient with respect to"
                        " the scanner coordinate frame of reference.");
 
+const OptionGroup KSpaceOption = OptionGroup ("Options for k-space filtering")
+  + Option ("window", "specify the shape of the k-space window filter; "
+                      "options are: " + join(Filter::kspace_window_choices, ",") + " "
+                      "(no default; must be specified for \"kspace\" operation)")
+    + Argument("name").type_choice(Filter::kspace_window_choices)
+  + Option ("strength", "modulate the strength of the chosen filter "
+                        "(exact interpretation & defaultmay depend on the exact filter chosen)")
+    + Argument("value").type_float(0.0, 1.0);
+
 const OptionGroup MedianOption = OptionGroup ("Options for median filter")
   + Option ("extent", "specify extent of median filtering neighbourhood in voxels."
                       " This can be specified either as a single value to be used for all 3 axes,"
@@ -134,6 +149,7 @@ void usage() {
   + FFTAxesOption
   + FFTOption
   + GradientOption
+  + KSpaceOption
   + MedianOption
   + NormaliseOption
   + SmoothOption
@@ -143,6 +159,7 @@ void usage() {
 }
 // clang-format on
 
+// TODO Use presence of SliceEncodingDirection to select defaults
 std::vector<size_t> get_axes(const Header &H, const std::vector<size_t> &default_axes) {
   auto opt = get_options("axes");
   std::vector<size_t> axes = default_axes;
@@ -172,7 +189,7 @@ void run() {
 
     const std::vector<size_t> inner_axes = get_axes(H_in, {0, 1});
 
-    Filter::Demodulate filter(input, inner_axes);
+    Filter::Demodulate filter(input, inner_axes, !get_options("linear").empty());
 
     Header H_out(H_in);
     Stride::set_from_command_line(H_out);
@@ -254,8 +271,42 @@ void run() {
     break;
   }
 
-  // Median
+  // k-space filtering
   case 3: {
+    auto opt_window = get_options("window");
+    if (opt_window.empty())
+      throw Exception("-window option is compulsory for k-space filtering");
+
+    Header H_in = Header::open(argument[0]);
+    const std::vector<size_t> axes = get_axes(H_in, {0, 1, 2});
+    const bool is_complex = H_in.datatype().is_complex();
+    auto input = H_in.get_image<cdouble>();
+
+    Image<double> window;
+    switch (Filter::kspace_windowfn_t(int(opt_window[0][0]))) {
+    case Filter::kspace_windowfn_t::TUKEY:
+      window = Filter::KSpace::window_tukey(H_in, axes, get_option_value("strength", Filter::default_tukey_width));
+      break;
+    default:
+      assert(false);
+    }
+    Filter::KSpace filter(H_in, window);
+    Header H_out(H_in);
+
+    if (is_complex) {
+      auto output = Image<cdouble>::create(argument[2], H_out);
+      filter(input, output);
+    } else {
+      H_out.datatype() = DataType::Float32;
+      H_out.datatype().set_byte_order_native();
+      auto output = Image<float>::create(argument[2], H_out);
+      filter(input, output);
+    }
+    break;
+  }
+
+  // Median
+  case 4: {
     auto input = Image<float>::open(argument[0]);
     Filter::Median filter(input);
 
@@ -272,7 +323,7 @@ void run() {
   }
 
   // Smooth
-  case 4: {
+  case 5: {
     auto input = Image<float>::open(argument[0]);
     Filter::Smooth filter(input);
 
@@ -303,7 +354,7 @@ void run() {
   }
 
   // Normalisation
-  case 5: {
+  case 6: {
     auto input = Image<float>::open(argument[0]);
     Filter::Normalise filter(input);
 
@@ -320,7 +371,7 @@ void run() {
   }
 
   // Zclean
-  case 6: {
+  case 7: {
     auto input = Image<float>::open(argument[0]);
     Filter::ZClean filter(input);
 
diff --git a/core/filter/demodulate.h b/core/filter/demodulate.h
index d28aad4741..1fc16e55c7 100644
--- a/core/filter/demodulate.h
+++ b/core/filter/demodulate.h
@@ -21,6 +21,7 @@
 #include "algo/copy.h"
 #include "algo/loop.h"
 #include "filter/base.h"
+#include "filter/kspace.h"
 #include "filter/smooth.h"
 #include "image.h"
 #include "interp/cubic.h"
@@ -31,13 +32,21 @@ namespace MR::Filter {
 /** \addtogroup Filters
 @{ */
 
+// From Manzano-Patron et al. 2024
+constexpr default_type default_tukey_FWHM_demodulate = 0.58;
+constexpr default_type default_tukey_alpha_demodulate = 2.0 * (1.0 - default_tukey_FWHM_demodulate);
+
 /*! Estimate a linear phase ramp of a complex image and demodulate by such
  */
 class Demodulate : public Base {
 public:
   template <class ImageType>
-  Demodulate(ImageType &in, const std::vector<size_t> &inner_axes)
-      : Base(in), phase(Image<cfloat>::scratch(in, "Scratch image storing linear phase for demodulator")) {
+  Demodulate(ImageType &in, const std::vector<size_t> &inner_axes, const bool linear)
+      : Base(in),
+        phase(Image<cfloat>::scratch(in,
+                                     std::string("Scratch image storing ")   //
+                                         + (linear ? "linear" : "nonlinear") //
+                                         + " phase for demodulator")) {      //
 
     using value_type = typename ImageType::value_type;
     using real_type = typename ImageType::value_type::value_type;
@@ -55,7 +64,8 @@ class Demodulate : public Base {
         outer_axes.push_back(axis);
     }
 
-    ProgressBar progress("estimating linear phase modulation", inner_axes.size() + 1);
+    ProgressBar progress(std::string("estimating ") + (linear ? "linear" : "nonlinear") + " phase modulation",
+                         inner_axes.size() + 1 + (linear ? 0 : inner_axes.size()));
 
     // FFT currently hard-wired to double precision;
     //   have to manually load into cdouble memory
@@ -83,12 +93,15 @@ class Demodulate : public Base {
           std::swap(temp, kspace);
           break;
         }
-        Math::FFT(temp, kspace, inner_axes[n], FFTW_FORWARD, true);
+        // Centred FFT if linear (so we can do peak-finding without wraparound);
+        //   not centred if non-linear (for compatibility with window filter generation function)
+        Math::FFT(temp, kspace, inner_axes[n], FFTW_FORWARD, linear);
         ++progress;
       }
     }
 
-    auto gen_phase =
+    // TODO Can likely remove "input" from here
+    auto gen_linear_phase =
         [&](Image<value_type> &input, Image<cdouble> &kspace, Image<cfloat> &phase, const std::vector<size_t> &axes) {
           std::array<bool, 3> axis_mask({false, false, false});
           for (auto axis : axes) {
@@ -171,13 +184,33 @@ class Demodulate : public Base {
           }
         };
 
-    if (outer_axes.size()) {
-      // TODO Multi-thread
-      // ThreadedLoop("Estimating phase ramps", input, outer_axes).run(gen_phase, input, kspace);
-      for (auto l_outer = Loop(outer_axes)(input, kspace, phase); l_outer; ++l_outer)
-        gen_phase(input, kspace, phase, inner_axes);
+    auto gen_nonlinear_phase =
+        [&](Image<value_type> &input, Image<cdouble> &kspace, Image<cfloat> &phase, const std::vector<size_t> &axes) {
+          Image<double> window = Filter::KSpace::window_tukey(input, axes, default_tukey_alpha_demodulate);
+          Adapter::Replicate<Image<double>> replicating_window(window, in);
+          for (auto l = Loop(kspace)(kspace, replicating_window); l; ++l)
+            kspace.value() *= double(replicating_window.value());
+          for (auto axis : axes) {
+            Math::FFT(kspace, kspace, axis, FFTW_BACKWARD, false);
+            ++progress;
+          }
+          for (auto l = Loop(kspace)(kspace, phase); l; ++l) {
+            cdouble value = cdouble(kspace.value());
+            value /= std::sqrt(norm(value));
+            phase.value() = {float(value.real()), float(value.imag())};
+          }
+        };
+
+    if (linear) {
+      if (outer_axes.size()) {
+        // TODO Multi-thread
+        for (auto l_outer = Loop(outer_axes)(input, kspace, phase); l_outer; ++l_outer)
+          gen_linear_phase(input, kspace, phase, inner_axes);
+      } else {
+        gen_linear_phase(input, kspace, phase, inner_axes);
+      }
     } else {
-      gen_phase(input, kspace, phase, inner_axes);
+      gen_nonlinear_phase(input, kspace, phase, inner_axes);
     }
   }
 
diff --git a/core/filter/kspace.h b/core/filter/kspace.h
new file mode 100644
index 0000000000..e6e99d3c45
--- /dev/null
+++ b/core/filter/kspace.h
@@ -0,0 +1,200 @@
+/* Copyright (c) 2008-2024 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#pragma once
+
+#include <string>
+#include <vector>
+
+#include "adapter/replicate.h"
+#include "filter/base.h"
+#include "header.h"
+#include "image.h"
+#include "math/fft.h"
+#include "types.h"
+
+namespace MR::Filter {
+
+std::vector<std::string> kspace_window_choices({"tukey"});
+enum class kspace_windowfn_t { TUKEY };
+constexpr default_type default_tukey_width = 0.5;
+
+/** \addtogroup Filters
+@{ */
+
+/*! Apply (or reverse) k-space filtering
+ */
+// TODO Add ability to undo prior filtering
+class KSpace : public Base {
+public:
+  KSpace(const Header &H, Image<double> &window) : Base(H), window(window) {
+    for (size_t axis = 0; axis != H.ndim(); ++axis) {
+      if (axis < window.ndim() && window.size(axis) > 1)
+        inner_axes.push_back(axis);
+      else
+        outer_axes.push_back(axis);
+    }
+  }
+
+  template <class InputImageType, class OutputImageType>
+  typename std::enable_if<std::is_same<typename InputImageType::value_type, std::complex<double>>::value, void>::type
+  operator()(InputImageType &in, OutputImageType &out) {
+    Image<cdouble> kspace;
+    Image<cdouble> temp;
+    {
+      for (ssize_t n = 0; n != inner_axes.size(); ++n) {
+        switch (n) {
+        case 0:
+          kspace = Image<cdouble>::scratch(in,
+                                           "Scratch k-space for \"" + in.name() + "\""        //
+                                               + (inner_axes.size() > 1 ? " (1 of 2)" : "")); //
+          Math::FFT(in, kspace, inner_axes[0], FFTW_FORWARD, false);
+          break;
+        case 1:
+          temp = kspace;
+          kspace = Image<cdouble>::scratch(in,
+                                           "Scratch k-space for \"" + in.name() + "\" " //
+                                               + "(2 of 2)");                           //
+          Math::FFT(temp, kspace, inner_axes[n], FFTW_FORWARD, false);
+          break;
+        default:
+          std::swap(temp, kspace);
+          Math::FFT(temp, kspace, inner_axes[n], FFTW_FORWARD, false);
+          break;
+        }
+      }
+    }
+
+    if (outer_axes.size()) {
+      Adapter::Replicate<Image<double>> replicating_window(window, in);
+      for (auto l = Loop(in)(kspace, replicating_window); l; ++l)
+        kspace.value() *= double(replicating_window.value());
+    } else {
+      for (auto l = Loop(in)(kspace, window); l; ++l)
+        kspace.value() *= double(window.value());
+    }
+
+    for (ssize_t n = 0; n != inner_axes.size(); ++n) {
+      if (n == inner_axes.size() - 1) {
+        // Final FFT:
+        //   use output image if applicable;
+        //   perform amplitude transform if necessary
+        do_final_fft(out, kspace, temp, inner_axes[n]);
+      } else {
+        Math::FFT(kspace, temp, inner_axes[n], FFTW_BACKWARD, false);
+        std::swap(kspace, temp);
+      }
+    }
+  }
+
+  template <class InputImageType, class OutputImageType>
+  typename std::enable_if<!std::is_same<typename InputImageType::value_type, cdouble>::value, void>::type
+  operator()(InputImageType &in, OutputImageType &out) {
+    Image<cdouble> temp = Image<cdouble>::scratch(in, "Scratch \"" + in.name() + "\" converted to cdouble for FFT");
+    for (auto l = Loop(in)(in, temp); l; ++l)
+      temp.value() = {double(in.value()), double(0)};
+    (*this)(temp, out);
+  }
+
+  // TODO Static functions to create different windows
+  //   (different windows may require different input parameters)
+  // TODO Currently extending 1D window functions to ND via outer product;
+  //   this is however not a unique choice
+  static Image<double> window_tukey(const Header &header,                  //
+                                    const std::vector<size_t> &inner_axes, //
+                                    const default_type cosine_frac) {      //
+    assert(cosine_frac >= 0.0 && cosine_frac <= 1.0);
+    Image<double> window = Image<double>::scratch(make_window_header(header, inner_axes),
+                                                  "Scratch Tukey filter window with alpha=" + str(cosine_frac));
+    for (auto l = Loop(window)(window); l; ++l)
+      window.value() = 1.0;
+    for (auto axis : inner_axes) {
+      const size_t N = header.size(axis);
+      Eigen::Array<double, Eigen::Dynamic, 1> window1d = Eigen::Array<double, Eigen::Dynamic, 1>::Ones(N);
+      const default_type transition_lower = 0.5 - 0.5 * cosine_frac;
+      const default_type transition_upper = 0.5 + 0.5 * cosine_frac;
+      // Beware of FFT being non-centred
+      for (size_t n = 0; n != N; ++n) {
+        const default_type pos = default_type(n) / default_type(N);
+        if (pos > transition_lower && pos < transition_upper)
+          window1d[n] = 0.5 + 0.5 * std::cos(2.0 * Math::pi * (pos - transition_lower) / cosine_frac);
+      }
+      window1d *= 1.0 / double(N);
+      // Need to loop over all inner axes other than the current one
+      std::vector<size_t> inner_excluding_axis;
+      for (auto a : inner_axes) {
+        if (a != axis)
+          inner_excluding_axis.push_back(a);
+      }
+      window.reset();
+      for (auto l = Loop(inner_excluding_axis)(window); l; ++l) {
+        for (size_t n = 0; n != N; ++n) {
+          window.index(axis) = n;
+          window.value() *= window1d[n];
+        }
+      }
+    }
+    return window;
+  }
+
+protected:
+  Image<double> window;
+  std::vector<size_t> inner_axes;
+  std::vector<size_t> outer_axes;
+
+  template <class ImageType>
+  typename std::enable_if<std::is_same<typename ImageType::value_type, cdouble>::value, void>::type
+  do_final_fft(ImageType &out, Image<cdouble> &kspace, Image<cdouble> &scratch, const size_t axis) {
+    TRACE;
+    Math::FFT(kspace, out, axis, FFTW_BACKWARD, false);
+  }
+  template <class ImageType>
+  typename std::enable_if<std::is_same<typename ImageType::value_type, cfloat>::value, void>::type
+  do_final_fft(ImageType &out, Image<cdouble> &kspace, Image<cdouble> &scratch, const size_t axis) {
+    TRACE;
+    assert(scratch.valid());
+    Math::FFT(kspace, scratch, axis, FFTW_BACKWARD, false);
+    for (auto l = Loop(out)(scratch, out); l; ++l)
+      out.value() = {float(cdouble(scratch.value()).real()), float(cdouble(scratch.value()).imag())};
+  }
+  template <class ImageType>
+  typename std::enable_if<!MR::is_complex<typename ImageType::value_type>::value, void>::type
+  do_final_fft(ImageType &out, Image<cdouble> &kspace, Image<cdouble> &scratch, const size_t axis) {
+    TRACE;
+    assert(scratch.valid());
+    Math::FFT(kspace, scratch, axis, FFTW_BACKWARD, false);
+    for (auto l = Loop(out)(scratch, out); l; ++l)
+      out.value() = std::abs(cdouble(scratch.value()));
+  }
+
+  static Header make_window_header(const Header &header, const std::vector<size_t> &inner_axes) {
+    Header H(header);
+    H.datatype() = DataType::Float64;
+    H.datatype().set_byte_order_native();
+    std::vector<size_t>::const_iterator it = inner_axes.begin();
+    for (size_t axis = 0; axis != header.ndim(); ++axis) {
+      if (it != inner_axes.end() && *it == axis)
+        ++it;
+      else
+        H.size(axis) = 1;
+    }
+    squeeze_dim(H);
+    return H;
+  }
+};
+//! @}
+
+} // namespace MR::Filter
diff --git a/src/denoise/demodulate.cpp b/src/denoise/demodulate.cpp
index 8719a391d7..f5c5a5e009 100644
--- a/src/denoise/demodulate.cpp
+++ b/src/denoise/demodulate.cpp
@@ -23,66 +23,85 @@ using namespace MR::App;
 
 namespace MR::Denoise {
 
-const char *const demodulation_description = "If the input data are of complex type, "
-                                             "then a linear phase term will be removed from each k-space prior to PCA. "
-                                             "In the absence of metadata indicating otherwise, "
-                                             "it is inferred that the first two axes correspond to acquired slices, "
-                                             "and different slices / volumes will be demodulated individually; "
-                                             "this behaviour can be modified using the -demod_axes option.";
+const char *const demodulation_description =
+    "If the input data are of complex type, "
+    "then a smooth non-linear phase will be demodulated removed from each k-space prior to PCA. "
+    "In the absence of metadata indicating otherwise, "
+    "it is inferred that the first two axes correspond to acquired slices, "
+    "and different slices / volumes will be demodulated individually; "
+    "this behaviour can be modified using the -demod_axes option. "
+    "A strictly linear phase term can instead be regressed from each k-space, "
+    "similarly to performed in Cordero-Grande et al. 2019, "
+    "by specifying -demodulate linear.";
 
-const OptionGroup demodulation_options =
-    OptionGroup("Options for phase demodulation of complex data") +
-    Option("nodemod", "disable phase demodulation")
-    // TODO Consider option to disable the remodulation of the output denoised series;
-    //   would need to turn this into a function call,
-    //   as that option would need to be omitted from dwi2noise
-    // Perhaps -nodemod, -noremod could be combined into a type_choice()?
-    // This wouldn't be able to also cover the future prospect of linear vs. non-linear phase demodulation;
-    //   maybe input phase demodulation being none / linear / nonlinear would be the better type_choice()?
-    +
-    Option("demod_axes", "comma-separated list of axis indices along which FFT can be applied for phase demodulation") +
-    Argument("axes").type_sequence_int();
+// clang-format off
+const OptionGroup demodulation_options = OptionGroup("Options for phase demodulation of complex data")
+  + Option("demodulate",
+           "select form of phase demodulation; "
+           "options are: " + join(demodulation_choices, ",") + " "
+           "(default: nonlinear)")
+    + Argument("mode").type_choice(demodulation_choices)
+  + Option("demod_axes",
+           "comma-separated list of axis indices along which FFT can be applied for phase demodulation")
+    + Argument("axes").type_sequence_int();
+// clang-format on
 
-std::vector<size_t> get_demodulation_axes(const Header &H) {
+Demodulation get_demodulation(const Header &H) {
   const bool complex = H.datatype().is_complex();
-  auto opt = App::get_options("nodemod");
-  if (!opt.empty()) {
-    if (!App::get_options("demod_axes").empty())
-      throw Exception("Options -nodemod and -demod_axes are mutually exclusive");
-    return std::vector<size_t>();
-  }
-  opt = App::get_options("demod_axes");
-  if (opt.empty()) {
+  auto opt_mode = get_options("demodulate");
+  auto opt_axes = get_options("demod_axes");
+  Demodulation result;
+  if (opt_mode.empty()) {
     if (complex) {
-      auto slice_encoding_it = H.keyval().find("SliceEncodingDirection");
-      if (slice_encoding_it == H.keyval().end()) {
-        INFO("No header information on slice encoding; assuming first two axes are within-slice");
-        return {0, 1};
-      } else {
-        auto dir = Axes::id2dir(slice_encoding_it->second);
-        std::vector<size_t> result;
-        for (size_t axis = 0; axis != 3; ++axis) {
-          if (!dir[axis])
-            result.push_back(axis);
-        }
-        INFO("For header SliceEncodingDirection=\"" + slice_encoding_it->second +
-             "\", "
-             "chose demodulation axes: " +
-             join(result, ","));
-        return result;
+      result.mode = demodulation_t::NONLINEAR;
+    } else {
+      if (!opt_axes.empty()) {
+        throw Exception("Option -demod_axes cannot be specified: "
+                        "no phase demodulation of magnitude data");
       }
     }
   } else {
-    if (!complex)
-      throw Exception("Cannot perform phase demodulation on magnitude input image");
-    auto result = parse_ints<size_t>(opt[0][0]);
-    for (auto axis : result) {
+    result.mode = demodulation_t(int(opt_mode[0][0]));
+    if (!complex) {
+      switch (result.mode) {
+      case demodulation_t::NONE:
+        WARN("Specifying -demodulate none is redundant: "
+             "never any phase demodulation for magnitude input data");
+        break;
+      default:
+        throw Exception("Phase modulation cannot be utilised for magnitude-only input data");
+      }
+    }
+  }
+  if (!complex)
+    return result;
+  if (opt_axes.empty()) {
+    auto slice_encoding_it = H.keyval().find("SliceEncodingDirection");
+    if (slice_encoding_it == H.keyval().end()) {
+      // TODO Ideally this would be the first two axes *on disk*,
+      //   not following transform realignment
+      INFO("No header information on slice encoding; "
+           "assuming first two axes are within-slice");
+      result.axes = {0, 1};
+    } else {
+      auto dir = Axes::id2dir(slice_encoding_it->second);
+      for (size_t axis = 0; axis != 3; ++axis) {
+        if (!dir[axis])
+          result.axes.push_back(axis);
+      }
+      INFO("For header SliceEncodingDirection=\"" + slice_encoding_it->second +
+           "\", "
+           "chose demodulation axes: " +
+           join(result.axes, ","));
+    }
+  } else {
+    result.axes = parse_ints<size_t>(opt_axes[0][0]);
+    for (auto axis : result.axes) {
       if (axis > 2)
         throw Exception("Phase demodulation implementation not yet robust to non-spatial axes");
     }
-    return result;
   }
-  return std::vector<size_t>();
+  return result;
 }
 
 } // namespace MR::Denoise
diff --git a/src/denoise/demodulate.h b/src/denoise/demodulate.h
index d55cc8078f..b0d9d04153 100644
--- a/src/denoise/demodulate.h
+++ b/src/denoise/demodulate.h
@@ -16,6 +16,7 @@
 
 #pragma once
 
+#include <string>
 #include <vector>
 
 #include "app.h"
@@ -25,8 +26,21 @@ namespace MR::Denoise {
 
 extern const char *const demodulation_description;
 
+const std::vector<std::string> demodulation_choices({"none", "linear", "nonlinear"});
+enum class demodulation_t { NONE, LINEAR, NONLINEAR };
+
 extern const App::OptionGroup demodulation_options;
 
-std::vector<size_t> get_demodulation_axes(const Header &);
+class Demodulation {
+public:
+  Demodulation(demodulation_t mode) : mode(mode) {}
+  Demodulation() : mode(demodulation_t::NONE) {}
+  explicit operator bool() const { return mode != demodulation_t::NONE; }
+  bool operator!() const { return mode == demodulation_t::NONE; }
+  demodulation_t mode;
+  std::vector<size_t> axes;
+};
+
+Demodulation get_demodulation(const Header &);
 
 } // namespace MR::Denoise

From 125ea3df37829fb26da73d19f9d3b0a589440aab Mon Sep 17 00:00:00 2001
From: Robert Smith <robert.smith@florey.edu.au>
Date: Fri, 6 Dec 2024 22:12:41 +1100
Subject: [PATCH 29/34] dwidenoise: New option -fixed_rank

Results in the output being reconstructed from a fixed number of eigenvectors, without making any attempt at estimating the signal rank / noise level from the data.
Closes 3046.
---
 cmd/dwi2noise.cpp                   |  2 +-
 cmd/dwidenoise.cpp                  |  8 ++---
 src/denoise/estimator/estimator.cpp | 47 +++++++++++++++++++++-------
 src/denoise/estimator/estimator.h   |  5 +--
 src/denoise/estimator/import.h      |  2 --
 src/denoise/estimator/rank.h        | 48 +++++++++++++++++++++++++++++
 6 files changed, 91 insertions(+), 21 deletions(-)
 create mode 100644 src/denoise/estimator/rank.h

diff --git a/cmd/dwi2noise.cpp b/cmd/dwi2noise.cpp
index 6c2eb17e62..ea8a078e5b 100644
--- a/cmd/dwi2noise.cpp
+++ b/cmd/dwi2noise.cpp
@@ -101,7 +101,7 @@ void usage() {
   OPTIONS
   + OptionGroup("Options for modifying PCA computations")
   + datatype_option
-  + Estimator::option
+  + Estimator::estimator_option
   + Kernel::options
   + subsample_option
   + demodulation_options
diff --git a/cmd/dwidenoise.cpp b/cmd/dwidenoise.cpp
index de9c31dd79..2f6f950ddc 100644
--- a/cmd/dwidenoise.cpp
+++ b/cmd/dwidenoise.cpp
@@ -31,6 +31,7 @@
 #include "denoise/estimator/estimator.h"
 #include "denoise/estimator/exp.h"
 #include "denoise/estimator/mrm2022.h"
+#include "denoise/estimator/rank.h"
 #include "denoise/estimator/result.h"
 #include "denoise/exports.h"
 #include "denoise/kernel/cuboid.h"
@@ -142,7 +143,7 @@ void usage() {
   OPTIONS
   + OptionGroup("Options for modifying PCA computations")
   + datatype_option
-  + Estimator::option
+  + Estimator::estimator_denoise_options
   + Kernel::options
   + subsample_option
   + demodulation_options
@@ -154,9 +155,6 @@ void usage() {
            "if noise level estimate is to be used for denoising also "
            "it must be additionally provided via the -noise_in option")
     + Argument("image").type_image_in()
-  + Option("noise_in",
-           "import a pre-estimated noise level map for noise removal rather than estimating this level from data")
-    + Argument("image").type_image_in()
 
   + OptionGroup("Options that affect reconstruction of the output image series")
   + Option("filter",
@@ -339,7 +337,7 @@ void run() {
   auto estimator = Estimator::make_estimator(true);
   assert(estimator);
 
-  filter_type filter = filter_type::OPTSHRINK;
+  filter_type filter = get_options("fixed_rank").empty() ? filter_type::OPTSHRINK : filter_type::TRUNCATE;
   opt = get_options("filter");
   if (!opt.empty())
     filter = filter_type(int(opt[0][0]));
diff --git a/src/denoise/estimator/estimator.cpp b/src/denoise/estimator/estimator.cpp
index 7720bf5254..41ffee1406 100644
--- a/src/denoise/estimator/estimator.cpp
+++ b/src/denoise/estimator/estimator.cpp
@@ -21,30 +21,55 @@
 #include "denoise/estimator/import.h"
 #include "denoise/estimator/med.h"
 #include "denoise/estimator/mrm2022.h"
+#include "denoise/estimator/rank.h"
 
 namespace MR::Denoise::Estimator {
 
 using namespace App;
 
-const Option option = Option("estimator",
-                             "Select the noise level estimator"
-                             " (default = Exp2),"
-                             " either: \n"
-                             "* Exp1: the original estimator used in Veraart et al. (2016); \n"
-                             "* Exp2: the improved estimator introduced in Cordero-Grande et al. (2019); \n"
-                             "* Med: estimate based on the median eigenvalue as in Gavish and Donohue (2014); \n"
-                             "* MRM2022: the alternative estimator introduced in Olesen et al. (2022).") +
-                      Argument("algorithm").type_choice(estimators);
+// clang-format off
+const Option estimator_option =
+    Option("estimator",
+           "Select the noise level estimator"
+           " (default = Exp2),"
+           " either: \n"
+           "* Exp1: the original estimator used in Veraart et al. (2016); \n"
+           "* Exp2: the improved estimator introduced in Cordero-Grande et al. (2019); \n"
+           "* Med: estimate based on the median eigenvalue as in Gavish and Donohue (2014); \n"
+           "* MRM2022: the alternative estimator introduced in Olesen et al. (2022). \n"
+           "Operation will be bypassed if -noise_in or -fixed_rank are specified")
+      + Argument("algorithm").type_choice(estimators);
 
-std::shared_ptr<Base> make_estimator(const bool permit_noise_in) {
+const OptionGroup estimator_denoise_options =
+    OptionGroup("Options relating to signal / noise level estimation for denoising")
+
+    + estimator_option
+
+    + Option("noise_in",
+             "import a pre-estimated noise level map for denoising rather than estimating this level from data")
+      + Argument("image").type_image_in()
+
+    + Option("fixed_rank",
+             "set a fixed input signal rank rather than estimating the noise level from the data")
+      + Argument("value").type_integer(1);
+
+std::shared_ptr<Base> make_estimator(const bool permit_bypass) {
   auto opt = get_options("estimator");
-  if (permit_noise_in) {
+  if (permit_bypass) {
     auto noise_in = get_options("noise_in");
+    auto fixed_rank = get_options("fixed_rank");
     if (!noise_in.empty()) {
       if (!opt.empty())
         throw Exception("Cannot both provide an input noise level image and specify a noise level estimator");
+      if (!fixed_rank.empty())
+        throw Exception("Cannot both provide an input noise level image and request a fixed signal rank");
       return std::make_shared<Import>(noise_in[0][0]);
     }
+    if (!fixed_rank.empty()) {
+      if (!opt.empty())
+        throw Exception("Cannot both provide an input signal rank and specify a noise level estimator");
+      return std::make_shared<Rank>(fixed_rank[0][0]);
+    }
   }
   const estimator_type est = opt.empty() ? estimator_type::EXP2 : estimator_type((int)(opt[0][0]));
   switch (est) {
diff --git a/src/denoise/estimator/estimator.h b/src/denoise/estimator/estimator.h
index 0a1cac5c8e..322d65ca96 100644
--- a/src/denoise/estimator/estimator.h
+++ b/src/denoise/estimator/estimator.h
@@ -26,9 +26,10 @@ namespace MR::Denoise::Estimator {
 
 class Base;
 
-extern const App::Option option;
+extern const App::Option estimator_option;
+extern const App::OptionGroup estimator_denoise_options;
 const std::vector<std::string> estimators = {"exp1", "exp2", "med", "mrm2022"};
 enum class estimator_type { EXP1, EXP2, MED, MRM2022 };
-std::shared_ptr<Base> make_estimator(const bool permit_noise_in);
+std::shared_ptr<Base> make_estimator(const bool permit_bypass);
 
 } // namespace MR::Denoise::Estimator
diff --git a/src/denoise/estimator/import.h b/src/denoise/estimator/import.h
index 546a6c1cd1..6705c11155 100644
--- a/src/denoise/estimator/import.h
+++ b/src/denoise/estimator/import.h
@@ -52,7 +52,6 @@ class Import : public Base {
     //   estimate the upper bound of the MP distribution and rank of signal
     //   given the ordered list of eigenvalues
     double cumulative_lambda = 0.0;
-    double recalc_sigmasq = 0.0;
     for (ssize_t p = 0; p != r; ++p) {
       const double lambda = std::max(s[p], 0.0) / q;
       cumulative_lambda += lambda;
@@ -60,7 +59,6 @@ class Import : public Base {
       if (sigma_sq < result.sigma2) {
         result.cutoff_p = p;
         result.lamplus = lambda;
-        recalc_sigmasq = sigma_sq;
       }
     }
     // TODO It would be nice if the upper bound, lambda_plus,
diff --git a/src/denoise/estimator/rank.h b/src/denoise/estimator/rank.h
new file mode 100644
index 0000000000..69a73f36ca
--- /dev/null
+++ b/src/denoise/estimator/rank.h
@@ -0,0 +1,48 @@
+/* Copyright (c) 2008-2024 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#pragma once
+
+#include "denoise/estimator/base.h"
+#include "denoise/estimator/result.h"
+
+namespace MR::Denoise::Estimator {
+
+class Rank : public Base {
+public:
+  Rank(const ssize_t r) : rank(r) {}
+  Result operator()(const eigenvalues_type &s,
+                    const ssize_t m,
+                    const ssize_t n,
+                    const Eigen::Vector3d & /*unused*/) const final {
+    Result result;
+    const ssize_t r = std::min(m, n);
+    const ssize_t q = std::max(m, n);
+    result.cutoff_p = r - rank;
+    double clam = 0.0;
+    for (ssize_t p = 0; p <= result.cutoff_p; ++p)
+      clam += std::max(s[p], 0.0);
+    clam /= q;
+    result.sigma2 = clam / (result.cutoff_p + 1);
+    result.lamplus = std::max(s[result.cutoff_p], 0.0);
+    return result;
+  }
+
+protected:
+  const ssize_t rank;
+};
+
+} // namespace MR::Denoise::Estimator

From 593bd53777837dad2a62fcd94b854cad81906f33 Mon Sep 17 00:00:00 2001
From: Robert Smith <robert.smith@florey.edu.au>
Date: Mon, 9 Dec 2024 15:42:33 +1100
Subject: [PATCH 30/34] dwidenoise: Fixes to nonstationarity correction

- Fix erroneous spatial localisation of voxels within patch for reading noise level estimate values.
- New option -noise_cov, for investigating the variance in within-patch noise level.
---
 cmd/dwi2noise.cpp        |  4 ++--
 cmd/dwidenoise.cpp       | 18 ++++++++++++++----
 src/denoise/estimate.cpp | 14 ++++++++++++--
 src/denoise/estimate.h   |  6 ++++++
 src/denoise/exports.h    |  2 ++
 5 files changed, 36 insertions(+), 8 deletions(-)

diff --git a/cmd/dwi2noise.cpp b/cmd/dwi2noise.cpp
index ea8a078e5b..ccff709750 100644
--- a/cmd/dwi2noise.cpp
+++ b/cmd/dwi2noise.cpp
@@ -231,12 +231,12 @@ void run() {
     prec += 2; // support complex input data
   switch (prec) {
   case 0:
-    assert(demodulation_axes.empty());
+    assert(demodulation.axes.empty());
     INFO("select real float32 for processing");
     run<float>(dwi, subsample, kernel, nonstationarity_image, estimator, exports);
     break;
   case 1:
-    assert(demodulation_axes.empty());
+    assert(demodulation.axes.empty());
     INFO("select real float64 for processing");
     run<double>(dwi, subsample, kernel, nonstationarity_image, estimator, exports);
     break;
diff --git a/cmd/dwidenoise.cpp b/cmd/dwidenoise.cpp
index 2f6f950ddc..3ffb47070a 100644
--- a/cmd/dwidenoise.cpp
+++ b/cmd/dwidenoise.cpp
@@ -147,8 +147,7 @@ void usage() {
   + Kernel::options
   + subsample_option
   + demodulation_options
-  // TODO If explicitly regressing the mean prior to Casorati formation,
-  //   this should happen _before_ rescaling based on noise level
+
   + Option("nonstationarity",
            "import an estimated map of noise nonstationarity; "
            "note that this will be used for within-patch non-stationariy correction only, "
@@ -202,6 +201,10 @@ void usage() {
   + Option("sum_optshrink",
            "the sum of eigenvector weights computed for the denoising patch centred at each voxel "
            "as a result of performing optimal shrinkage")
+    + Argument("image").type_image_out()
+    + Option("noise_cov",
+           "export an image of the Coefficient of Variation (CoV) of noise level within each patch "
+           "(only applicable if -nonstationarity is specified)")
     + Argument("image").type_image_out();
 
   COPYRIGHT =
@@ -397,17 +400,24 @@ void run() {
     exports.set_sum_aggregation("");
   }
 
+  opt = get_options("noise_cov");
+  if (!opt.empty()) {
+    if (!nonstationarity_image.valid())
+      throw Exception("-noise_variance can only be specified if -nonstationarity option is used");
+    exports.set_noise_cov(opt[0][0]);
+  }
+
   int prec = get_option_value("datatype", 0); // default: single precision
   if (dwi.datatype().is_complex())
     prec += 2; // support complex input data
   switch (prec) {
   case 0:
-    assert(demodulation_axes.empty());
+    assert(demodulation.axes.empty());
     INFO("select real float32 for processing");
     run<float>(dwi, subsample, kernel, nonstationarity_image, estimator, filter, aggregator, argument[1], exports);
     break;
   case 1:
-    assert(demodulation_axes.empty());
+    assert(demodulation.axes.empty());
     INFO("select real float64 for processing");
     run<double>(dwi, subsample, kernel, nonstationarity_image, estimator, filter, aggregator, argument[1], exports);
     break;
diff --git a/src/denoise/estimate.cpp b/src/denoise/estimate.cpp
index 0b6e7be5c1..bba04ee2b6 100644
--- a/src/denoise/estimate.cpp
+++ b/src/denoise/estimate.cpp
@@ -34,6 +34,7 @@ Estimate<F>::Estimate(const Header &header,
       subsample(subsample),
       kernel(kernel),
       estimator(estimator),
+      transform(std::make_shared<Transform>(header)),
       nonstationarity_image(nonstationarity_image),
       X(m, kernel->estimated_size()),
       XtX(std::min(m, kernel->estimated_size()), std::min(m, kernel->estimated_size())),
@@ -121,6 +122,14 @@ template <typename F> void Estimate<F>::operator()(Image<F> &dwi) {
       exports.patchcount.value() = exports.patchcount.value() + 1;
     }
   }
+  if (exports.noise_cov.valid()) {
+    double variance(double(0));
+    for (auto v : patch.voxels)
+      variance += Math::pow2(v.noise_level - patch.centre_noise);
+    variance /= (patch.voxels.size() - 1);
+    assign_pos_of(ss_index).to(exports.noise_cov);
+    exports.noise_cov.value() = std::sqrt(variance) / patch.centre_noise;
+  }
 }
 
 template <typename F> void Estimate<F>::load_data(Image<F> &image) {
@@ -132,15 +141,16 @@ template <typename F> void Estimate<F>::load_data(Image<F> &image) {
     assert(!(!interp));
     patch.centre_noise = interp.value();
     for (ssize_t i = 0; i != patch.voxels.size(); ++i) {
-      interp.scanner(image.transform() * patch.voxels[i].index.cast<default_type>());
+      interp.scanner(transform->voxel2scanner * patch.voxels[i].index.cast<default_type>());
       // TODO Trying to pull intensity information from voxels beyond the extremities of the subsampled image
       //   may cause problems
       assert(!(!interp));
       const double voxel_noise = interp.value();
       patch.voxels[i].noise_level = voxel_noise;
+      const double scaling_factor = patch.centre_noise / voxel_noise;
       assign_pos_of(patch.voxels[i].index, 0, 3).to(image);
       X.col(i) = image.row(3);
-      X.col(i) *= patch.centre_noise / voxel_noise;
+      X.col(i) *= scaling_factor;
     }
   } else {
     for (ssize_t i = 0; i != patch.voxels.size(); ++i) {
diff --git a/src/denoise/estimate.h b/src/denoise/estimate.h
index 82f0f6dcc8..099166be47 100644
--- a/src/denoise/estimate.h
+++ b/src/denoise/estimate.h
@@ -31,6 +31,7 @@
 #include "denoise/subsample.h"
 #include "header.h"
 #include "image.h"
+#include "transform.h"
 
 namespace MR::Denoise {
 
@@ -56,6 +57,9 @@ template <typename F> class Estimate {
   std::shared_ptr<Kernel::Base> kernel;
   std::shared_ptr<Estimator::Base> estimator;
 
+  // Necessary for transform from input voxel locations to nonstationarity image
+  std::shared_ptr<Transform> transform;
+
   // Reusable memory
   Kernel::Data patch;
   Image<float> nonstationarity_image;
@@ -66,6 +70,8 @@ template <typename F> class Estimate {
   Estimator::Result threshold;
 
   // Export images
+  // Note: One instance created per thread,
+  //   so that when possible output image data can be written without mutex-locking
   Exports exports;
 
   // Some data can only be written in a thread-safe manner
diff --git a/src/denoise/exports.h b/src/denoise/exports.h
index d93008025c..0cd1e8a715 100644
--- a/src/denoise/exports.h
+++ b/src/denoise/exports.h
@@ -59,6 +59,7 @@ class Exports {
     else
       sum_aggregation = Image<float>::create(path, H_in);
   }
+  void set_noise_cov(const std::string &path) { noise_cov = Image<float>::create(path, H_ss); }
 
   Image<float> noise_out;
   Image<uint16_t> rank_input;
@@ -68,6 +69,7 @@ class Exports {
   Image<uint16_t> voxelcount;
   Image<uint16_t> patchcount;
   Image<float> sum_aggregation;
+  Image<float> noise_cov;
 
 protected:
   Header H_in;

From 471be567d71e032bbb58dfb56dda6958533b0ddb Mon Sep 17 00:00:00 2001
From: Robert Smith <robert.smith@florey.edu.au>
Date: Mon, 9 Dec 2024 16:54:41 +1100
Subject: [PATCH 31/34] dwi*noise: Improve handling where estimator fails

---
 src/denoise/estimate.cpp      | 37 ++++++++++---------
 src/denoise/estimator/base.h  |  1 +
 src/denoise/estimator/exp.cpp | 68 +++++++++++++++++++++++++++++++++++
 src/denoise/estimator/exp.h   | 48 +++++++------------------
 src/denoise/recon.cpp         |  3 +-
 5 files changed, 105 insertions(+), 52 deletions(-)
 create mode 100644 src/denoise/estimator/exp.cpp

diff --git a/src/denoise/estimate.cpp b/src/denoise/estimate.cpp
index bba04ee2b6..2eb23c1168 100644
--- a/src/denoise/estimate.cpp
+++ b/src/denoise/estimate.cpp
@@ -84,6 +84,7 @@ template <typename F> void Estimate<F>::operator()(Image<F> &dwi) {
 #endif
 
   load_data(dwi);
+  assert(X.leftCols(n).allFinite());
 
   // Compute Eigendecomposition:
   if (m <= n)
@@ -140,24 +141,28 @@ template <typename F> void Estimate<F>::load_data(Image<F> &image) {
     interp.scanner(patch.centre_realspace);
     assert(!(!interp));
     patch.centre_noise = interp.value();
-    for (ssize_t i = 0; i != patch.voxels.size(); ++i) {
-      interp.scanner(transform->voxel2scanner * patch.voxels[i].index.cast<default_type>());
-      // TODO Trying to pull intensity information from voxels beyond the extremities of the subsampled image
-      //   may cause problems
-      assert(!(!interp));
-      const double voxel_noise = interp.value();
-      patch.voxels[i].noise_level = voxel_noise;
-      const double scaling_factor = patch.centre_noise / voxel_noise;
-      assign_pos_of(patch.voxels[i].index, 0, 3).to(image);
-      X.col(i) = image.row(3);
-      X.col(i) *= scaling_factor;
-    }
-  } else {
-    for (ssize_t i = 0; i != patch.voxels.size(); ++i) {
-      assign_pos_of(patch.voxels[i].index, 0, 3).to(image);
-      X.col(i) = image.row(3);
+    if (patch.centre_noise > 0.0) {
+      for (ssize_t i = 0; i != patch.voxels.size(); ++i) {
+        interp.scanner(transform->voxel2scanner * patch.voxels[i].index.cast<default_type>());
+        // TODO Trying to pull intensity information from voxels beyond the extremities of the subsampled image
+        //   may cause problems
+        assert(!(!interp));
+        const double voxel_noise = interp.value();
+        patch.voxels[i].noise_level = voxel_noise;
+        const double scaling_factor = voxel_noise > 0.0 ? (patch.centre_noise / voxel_noise) : 1.0;
+        assert(std::isfinite(scaling_factor));
+        assign_pos_of(patch.voxels[i].index, 0, 3).to(image);
+        X.col(i) = image.row(3);
+        X.col(i) *= scaling_factor;
+      }
+      assign_pos_of(pos, 0, 3).to(image);
+      return;
     }
   }
+  for (ssize_t i = 0; i != patch.voxels.size(); ++i) {
+    assign_pos_of(patch.voxels[i].index, 0, 3).to(image);
+    X.col(i) = image.row(3);
+  }
   assign_pos_of(pos, 0, 3).to(image);
 }
 
diff --git a/src/denoise/estimator/base.h b/src/denoise/estimator/base.h
index a95b85d980..efa6190a8f 100644
--- a/src/denoise/estimator/base.h
+++ b/src/denoise/estimator/base.h
@@ -24,6 +24,7 @@ namespace MR::Denoise::Estimator {
 class Base {
 public:
   Base() = default;
+  Base(const Base &) = delete;
   virtual Result operator()(const eigenvalues_type &eigenvalues,
                             const ssize_t m,
                             const ssize_t n,
diff --git a/src/denoise/estimator/exp.cpp b/src/denoise/estimator/exp.cpp
new file mode 100644
index 0000000000..b06f2a1a29
--- /dev/null
+++ b/src/denoise/estimator/exp.cpp
@@ -0,0 +1,68 @@
+/* Copyright (c) 2008-2024 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#include "denoise/estimator/exp.h"
+
+namespace MR::Denoise::Estimator {
+
+template <ssize_t version>
+Result Exp<version>::operator()(const eigenvalues_type &s,
+                                const ssize_t m,
+                                const ssize_t n,
+                                const Eigen::Vector3d & /*unused*/) const {
+  Result result;
+  const ssize_t r = std::min(m, n);
+  const ssize_t q = std::max(m, n);
+  const double lam_r = std::max(s[0], 0.0) / q;
+  double clam = 0.0;
+  for (ssize_t p = 0; p < r; ++p) // p+1 is the number of noise components
+  {                               // (as opposed to the paper where p is defined as the number of signal components)
+    const double lam = std::max(s[p], 0.0) / q;
+    clam += lam;
+    double denominator = std::numeric_limits<double>::signaling_NaN();
+    switch (version) {
+    case 1:
+      denominator = q;
+      break;
+    case 2:
+      denominator = q - (r - p - 1);
+      break;
+    default:
+      assert(false);
+    }
+    const double gam = double(p + 1) / denominator;
+    const double sigsq1 = clam / double(p + 1);
+    const double sigsq2 = (lam - lam_r) / (4.0 * std::sqrt(gam));
+    // sigsq2 > sigsq1 if signal else noise
+    if (sigsq2 < sigsq1) {
+      result.sigma2 = sigsq1;
+      result.cutoff_p = p + 1;
+      result.lamplus = lam;
+    }
+  }
+  if (result.cutoff_p == -1) {
+    result.cutoff_p = 0;
+    result.sigma2 = 0.0;
+    result.lamplus = 0.0;
+    failure_count.fetch_add(1);
+  }
+  return result;
+}
+
+template class Exp<1>;
+template class Exp<2>;
+
+} // namespace MR::Denoise::Estimator
diff --git a/src/denoise/estimator/exp.h b/src/denoise/estimator/exp.h
index b0100c72e5..4e62cd9dc9 100644
--- a/src/denoise/estimator/exp.h
+++ b/src/denoise/estimator/exp.h
@@ -16,51 +16,29 @@
 
 #pragma once
 
+#include <atomic>
+
 #include "denoise/estimator/base.h"
 #include "denoise/estimator/result.h"
 
 namespace MR::Denoise::Estimator {
 
-// TODO Move to .cpp
 template <ssize_t version> class Exp : public Base {
 public:
-  Exp() = default;
+  Exp() : failure_count(0) {}
+  ~Exp() {
+    const ssize_t total = failure_count.load();
+    if (total > 0) {
+      WARN("Noise level estimator failed to converge for " + str(total) + " patches");
+    }
+  }
   Result operator()(const eigenvalues_type &s,
                     const ssize_t m,
                     const ssize_t n,
-                    const Eigen::Vector3d & /*unused*/) const final {
-    Result result;
-    const ssize_t r = std::min(m, n);
-    const ssize_t q = std::max(m, n);
-    const double lam_r = std::max(s[0], 0.0) / q;
-    double clam = 0.0;
-    for (ssize_t p = 0; p < r; ++p) // p+1 is the number of noise components
-    {                               // (as opposed to the paper where p is defined as the number of signal components)
-      const double lam = std::max(s[p], 0.0) / q;
-      clam += lam;
-      double denominator = std::numeric_limits<double>::signaling_NaN();
-      switch (version) {
-      case 1:
-        denominator = q;
-        break;
-      case 2:
-        denominator = q - (r - p - 1);
-        break;
-      default:
-        assert(false);
-      }
-      const double gam = double(p + 1) / denominator;
-      const double sigsq1 = clam / double(p + 1);
-      const double sigsq2 = (lam - lam_r) / (4.0 * std::sqrt(gam));
-      // sigsq2 > sigsq1 if signal else noise
-      if (sigsq2 < sigsq1) {
-        result.sigma2 = sigsq1;
-        result.cutoff_p = p + 1;
-        result.lamplus = lam;
-      }
-    }
-    return result;
-  }
+                    const Eigen::Vector3d & /*unused*/) const final;
+
+protected:
+  mutable std::atomic<ssize_t> failure_count;
 };
 
 } // namespace MR::Denoise::Estimator
diff --git a/src/denoise/recon.cpp b/src/denoise/recon.cpp
index b428d94f3f..cdba213285 100644
--- a/src/denoise/recon.cpp
+++ b/src/denoise/recon.cpp
@@ -182,7 +182,8 @@ template <typename F> void Recon<F>::operator()(Image<F> &dwi, Image<F> &out) {
     // Undo prior within-patch non-stationarity correction
     if (std::isfinite(Estimate<F>::patch.centre_noise)) {
       for (ssize_t i = 0; i != n; ++i)
-        Xr.col(i) *= Estimate<F>::patch.voxels[i].noise_level / Estimate<F>::patch.centre_noise;
+        if (Estimate<F>::patch.voxels[i].noise_level > 0.0)
+          Xr.col(i) *= Estimate<F>::patch.voxels[i].noise_level / Estimate<F>::patch.centre_noise;
     }
     std::lock_guard<std::mutex> lock(Estimate<F>::mutex);
     for (size_t voxel_index = 0; voxel_index != Estimate<F>::patch.voxels.size(); ++voxel_index) {

From 79e9c6527e1381eb08421d43dafe2ed61512bf6f Mon Sep 17 00:00:00 2001
From: Robert Smith <robert.smith@florey.edu.au>
Date: Tue, 10 Dec 2024 15:30:05 +1100
Subject: [PATCH 32/34] dwi*noise: Rename -nonstationarity -> -vst

Defining this operation as a variance-stabilising transformation is more accurate, and also flags the fact that that future augmentation accounting for the Rician bias in magnitude data will fall under the same mechanism.
---
 cmd/dwi2noise.cpp        | 30 +++++++-------
 cmd/dwidenoise.cpp       | 86 +++++++++++++++++++++++++---------------
 src/denoise/estimate.cpp |  8 ++--
 src/denoise/estimate.h   |  4 +-
 src/denoise/recon.cpp    |  6 +--
 src/denoise/recon.h      |  2 +-
 6 files changed, 78 insertions(+), 58 deletions(-)

diff --git a/cmd/dwi2noise.cpp b/cmd/dwi2noise.cpp
index ccff709750..8e94586bc5 100644
--- a/cmd/dwi2noise.cpp
+++ b/cmd/dwi2noise.cpp
@@ -90,7 +90,7 @@ void usage() {
     "Magnetic Resonance in Medicine, 2022, 89(3), 1160-1172"
 
   + "* If using -estimator med: "
-    "Gavish, M.; Donoho, D.L."
+    "Gavish, M.; Donoho, D.L. "
     "The Optimal Hard Threshold for Singular Values is 4/sqrt(3). "
     "IEEE Transactions on Information Theory, 2014, 60(8), 5040-5053.";
 
@@ -105,8 +105,8 @@ void usage() {
   + Kernel::options
   + subsample_option
   + demodulation_options
-  + Option("nonstationarity",
-           "import an estimated map of noise nonstationarity; "
+  + Option("vst",
+           "apply a within-patch variance-stabilising transformation based on a pre-estimated noise level map; "
            "note that this will be used for within-patch non-stationariy correction only, "
            "the output noise level estimate will still be derived from the input data")
     + Argument("image").type_image_in()
@@ -159,11 +159,11 @@ template <typename T>
 void run(Header &data,
          std::shared_ptr<Subsample> subsample,
          std::shared_ptr<Kernel::Base> kernel,
-         Image<float> &nonstationarity_image,
+         Image<float> &vst_noise_image,
          std::shared_ptr<Estimator::Base> estimator,
          Exports &exports) {
   auto input = data.get_image<T>().with_direct_io(3);
-  Estimate<T> func(data, subsample, kernel, nonstationarity_image, estimator, exports);
+  Estimate<T> func(data, subsample, kernel, vst_noise_image, estimator, exports);
   ThreadedLoop("running MP-PCA noise level estimation", data, 0, 3).run(func, input);
 }
 
@@ -172,11 +172,11 @@ void run(Header &data,
          const Demodulation &demodulation,
          std::shared_ptr<Subsample> subsample,
          std::shared_ptr<Kernel::Base> kernel,
-         Image<float> &nonstationarity_image,
+         Image<float> &vst_noise_image,
          std::shared_ptr<Estimator::Base> estimator,
          Exports &exports) {
   if (!demodulation) {
-    run<T>(data, subsample, kernel, nonstationarity_image, estimator, exports);
+    run<T>(data, subsample, kernel, vst_noise_image, estimator, exports);
     return;
   }
   auto input = data.get_image<T>();
@@ -185,7 +185,7 @@ void run(Header &data,
     Filter::Demodulate demodulator(input, demodulation.axes, demodulation.mode == demodulation_t::LINEAR);
     demodulator(input, input_demod);
   }
-  Estimate<T> func(data, subsample, kernel, nonstationarity_image, estimator, exports);
+  Estimate<T> func(data, subsample, kernel, vst_noise_image, estimator, exports);
   ThreadedLoop("running MP-PCA noise level estimation", data, 0, 3).run(func, input_demod);
 }
 
@@ -203,10 +203,10 @@ void run() {
   auto kernel = Kernel::make_kernel(dwi, subsample->get_factors());
   assert(kernel);
 
-  Image<float> nonstationarity_image;
-  auto opt = get_options("nonstationarity");
+  Image<float> vst_noise_image;
+  auto opt = get_options("vst");
   if (!opt.empty())
-    nonstationarity_image = Image<float>::open(opt[0][0]);
+    vst_noise_image = Image<float>::open(opt[0][0]);
 
   auto estimator = Estimator::make_estimator(false);
   assert(estimator);
@@ -233,20 +233,20 @@ void run() {
   case 0:
     assert(demodulation.axes.empty());
     INFO("select real float32 for processing");
-    run<float>(dwi, subsample, kernel, nonstationarity_image, estimator, exports);
+    run<float>(dwi, subsample, kernel, vst_noise_image, estimator, exports);
     break;
   case 1:
     assert(demodulation.axes.empty());
     INFO("select real float64 for processing");
-    run<double>(dwi, subsample, kernel, nonstationarity_image, estimator, exports);
+    run<double>(dwi, subsample, kernel, vst_noise_image, estimator, exports);
     break;
   case 2:
     INFO("select complex float32 for processing");
-    run<cfloat>(dwi, demodulation, subsample, kernel, nonstationarity_image, estimator, exports);
+    run<cfloat>(dwi, demodulation, subsample, kernel, vst_noise_image, estimator, exports);
     break;
   case 3:
     INFO("select complex float64 for processing");
-    run<cdouble>(dwi, demodulation, subsample, kernel, nonstationarity_image, estimator, exports);
+    run<cdouble>(dwi, demodulation, subsample, kernel, vst_noise_image, estimator, exports);
     break;
   }
 }
diff --git a/cmd/dwidenoise.cpp b/cmd/dwidenoise.cpp
index 3ffb47070a..b2da3e4524 100644
--- a/cmd/dwidenoise.cpp
+++ b/cmd/dwidenoise.cpp
@@ -148,8 +148,8 @@ void usage() {
   + subsample_option
   + demodulation_options
 
-  + Option("nonstationarity",
-           "import an estimated map of noise nonstationarity; "
+  + Option("vst",
+           "apply a within-patch variance-stabilising transformation based on a pre-estimated noise level map; "
            "note that this will be used for within-patch non-stationariy correction only, "
            "if noise level estimate is to be used for denoising also "
            "it must be additionally provided via the -noise_in option")
@@ -245,7 +245,7 @@ template <typename T>
 void run(Header &data,
          std::shared_ptr<Subsample> subsample,
          std::shared_ptr<Kernel::Base> kernel,
-         Image<float> &nonstationarity_image,
+         Image<float> &vst_noise_image,
          std::shared_ptr<Estimator::Base> estimator,
          filter_type filter,
          aggregator_type aggregator,
@@ -257,7 +257,7 @@ void run(Header &data,
   header.datatype() = DataType::from<T>();
   auto output = Image<T>::create(output_name, header);
   // run
-  Recon<T> func(data, subsample, kernel, nonstationarity_image, estimator, filter, aggregator, exports);
+  Recon<T> func(data, subsample, kernel, vst_noise_image, estimator, filter, aggregator, exports);
   ThreadedLoop("running MP-PCA denoising", data, 0, 3).run(func, input, output);
   // Rescale output if aggregation was performed
   if (aggregator == aggregator_type::EXCLUSIVE)
@@ -277,14 +277,14 @@ void run(Header &data,
          const Demodulation &demodulation,
          std::shared_ptr<Subsample> subsample,
          std::shared_ptr<Kernel::Base> kernel,
-         Image<float> &nonstationarity_image,
+         Image<float> &vst_noise_image,
          std::shared_ptr<Estimator::Base> estimator,
          filter_type filter,
          aggregator_type aggregator,
          const std::string &output_name,
          Exports &exports) {
   if (!demodulation) {
-    run<T>(data, subsample, kernel, nonstationarity_image, estimator, filter, aggregator, output_name, exports);
+    run<T>(data, subsample, kernel, vst_noise_image, estimator, filter, aggregator, output_name, exports);
     return;
   }
   auto input = data.get_image<T>();
@@ -302,7 +302,7 @@ void run(Header &data,
   header.datatype() = DataType::from<T>();
   auto output = Image<T>::create(output_name, header);
   // run
-  Recon<T> func(data, subsample, kernel, nonstationarity_image, estimator, filter, aggregator, exports);
+  Recon<T> func(data, subsample, kernel, vst_noise_image, estimator, filter, aggregator, exports);
   ThreadedLoop("running MP-PCA denoising", data, 0, 3).run(func, input_demodulated, output);
   // Re-apply phase ramps that were previously demodulated
   demodulate(output, true);
@@ -332,10 +332,10 @@ void run() {
   auto kernel = Kernel::make_kernel(dwi, subsample->get_factors());
   assert(kernel);
 
-  Image<float> nonstationarity_image;
-  auto opt = get_options("nonstationarity");
+  Image<float> vst_noise_image;
+  auto opt = get_options("vst");
   if (!opt.empty())
-    nonstationarity_image = Image<float>::open(opt[0][0]);
+    vst_noise_image = Image<float>::open(opt[0][0]);
 
   auto estimator = Estimator::make_estimator(true);
   assert(estimator);
@@ -402,7 +402,7 @@ void run() {
 
   opt = get_options("noise_cov");
   if (!opt.empty()) {
-    if (!nonstationarity_image.valid())
+    if (!vst_noise_image.valid())
       throw Exception("-noise_variance can only be specified if -nonstationarity option is used");
     exports.set_noise_cov(opt[0][0]);
   }
@@ -414,38 +414,58 @@ void run() {
   case 0:
     assert(demodulation.axes.empty());
     INFO("select real float32 for processing");
-    run<float>(dwi, subsample, kernel, nonstationarity_image, estimator, filter, aggregator, argument[1], exports);
+    run<float>(          //
+        dwi,             //
+        subsample,       //
+        kernel,          //
+        vst_noise_image, //
+        estimator,       //
+        filter,          //
+        aggregator,      //
+        argument[1],     //
+        exports);        //
     break;
   case 1:
     assert(demodulation.axes.empty());
     INFO("select real float64 for processing");
-    run<double>(dwi, subsample, kernel, nonstationarity_image, estimator, filter, aggregator, argument[1], exports);
+    run<double>(         //
+        dwi,             //
+        subsample,       //
+        kernel,          //
+        vst_noise_image, //
+        estimator,       //
+        filter,          //
+        aggregator,      //
+        argument[1],     //
+        exports);        //
     break;
   case 2:
     INFO("select complex float32 for processing");
-    run<cfloat>(dwi,
-                demodulation,
-                subsample,
-                kernel,
-                nonstationarity_image,
-                estimator,
-                filter,
-                aggregator,
-                argument[1],
-                exports);
+    run<cfloat>(         //
+        dwi,             //
+        demodulation,    //
+        subsample,       //
+        kernel,          //
+        vst_noise_image, //
+        estimator,       //
+        filter,          //
+        aggregator,      //
+        argument[1],     //
+        exports);        //
     break;
   case 3:
     INFO("select complex float64 for processing");
-    run<cdouble>(dwi,
-                 demodulation,
-                 subsample,
-                 kernel,
-                 nonstationarity_image,
-                 estimator,
-                 filter,
-                 aggregator,
-                 argument[1],
-                 exports);
+    run<cdouble>(        //
+        dwi,             //
+        demodulation,    //
+        subsample,       //
+        kernel,          //
+        vst_noise_image, //
+        estimator,       //
+        filter,          //
+        aggregator,      //
+        argument[1],     //
+        exports);        //
     break;
   }
 }
diff --git a/src/denoise/estimate.cpp b/src/denoise/estimate.cpp
index 2eb23c1168..3fb8da2af7 100644
--- a/src/denoise/estimate.cpp
+++ b/src/denoise/estimate.cpp
@@ -27,7 +27,7 @@ template <typename F>
 Estimate<F>::Estimate(const Header &header,
                       std::shared_ptr<Subsample> subsample,
                       std::shared_ptr<Kernel::Base> kernel,
-                      Image<float> &nonstationarity_image,
+                      Image<float> &vst_noise_image,
                       std::shared_ptr<Estimator::Base> estimator,
                       Exports &exports)
     : m(header.size(3)),
@@ -35,7 +35,7 @@ Estimate<F>::Estimate(const Header &header,
       kernel(kernel),
       estimator(estimator),
       transform(std::make_shared<Transform>(header)),
-      nonstationarity_image(nonstationarity_image),
+      vst_noise_image(vst_noise_image),
       X(m, kernel->estimated_size()),
       XtX(std::min(m, kernel->estimated_size()), std::min(m, kernel->estimated_size())),
       eig(std::min(m, kernel->estimated_size())),
@@ -135,9 +135,9 @@ template <typename F> void Estimate<F>::operator()(Image<F> &dwi) {
 
 template <typename F> void Estimate<F>::load_data(Image<F> &image) {
   const Kernel::Voxel::index_type pos({image.index(0), image.index(1), image.index(2)});
-  if (nonstationarity_image.valid()) {
+  if (vst_noise_image.valid()) {
     assert(patch.centre_realspace.allFinite());
-    Interp::Cubic<Image<float>> interp(nonstationarity_image);
+    Interp::Cubic<Image<float>> interp(vst_noise_image);
     interp.scanner(patch.centre_realspace);
     assert(!(!interp));
     patch.centre_noise = interp.value();
diff --git a/src/denoise/estimate.h b/src/denoise/estimate.h
index 099166be47..622acdd7f7 100644
--- a/src/denoise/estimate.h
+++ b/src/denoise/estimate.h
@@ -43,7 +43,7 @@ template <typename F> class Estimate {
   Estimate(const Header &header,
            std::shared_ptr<Subsample> subsample,
            std::shared_ptr<Kernel::Base> kernel,
-           Image<float> &nonstationarity_image,
+           Image<float> &vst_noise_image,
            std::shared_ptr<Estimator::Base> estimator,
            Exports &exports);
 
@@ -62,7 +62,7 @@ template <typename F> class Estimate {
 
   // Reusable memory
   Kernel::Data patch;
-  Image<float> nonstationarity_image;
+  Image<float> vst_noise_image;
   MatrixType X;
   MatrixType XtX;
   Eigen::SelfAdjointEigenSolver<MatrixType> eig;
diff --git a/src/denoise/recon.cpp b/src/denoise/recon.cpp
index cdba213285..9cfc177667 100644
--- a/src/denoise/recon.cpp
+++ b/src/denoise/recon.cpp
@@ -24,12 +24,12 @@ template <typename F>
 Recon<F>::Recon(const Header &header,
                 std::shared_ptr<Subsample> subsample,
                 std::shared_ptr<Kernel::Base> kernel,
-                Image<float> &nonstationarity_image,
+                Image<float> &vst_noise_image,
                 std::shared_ptr<Estimator::Base> estimator,
                 filter_type filter,
                 aggregator_type aggregator,
                 Exports &exports)
-    : Estimate<F>(header, subsample, kernel, nonstationarity_image, estimator, exports),
+    : Estimate<F>(header, subsample, kernel, vst_noise_image, estimator, exports),
       filter(filter),
       aggregator(aggregator),
       // FWHM = 2 x cube root of spacings between kernels
@@ -179,7 +179,7 @@ template <typename F> void Recon<F>::operator()(Image<F> &dwi, Image<F> &out) {
             Estimate<F>::eig.eigenvectors().adjoint())); //
     }
     assert(Xr.leftCols(n).allFinite());
-    // Undo prior within-patch non-stationarity correction
+    // Undo prior within-patch variance-stabilising transform
     if (std::isfinite(Estimate<F>::patch.centre_noise)) {
       for (ssize_t i = 0; i != n; ++i)
         if (Estimate<F>::patch.voxels[i].noise_level > 0.0)
diff --git a/src/denoise/recon.h b/src/denoise/recon.h
index 977df221f4..3d96dedc3d 100644
--- a/src/denoise/recon.h
+++ b/src/denoise/recon.h
@@ -37,7 +37,7 @@ template <typename F> class Recon : public Estimate<F> {
   Recon(const Header &header,
         std::shared_ptr<Subsample> subsample,
         std::shared_ptr<Kernel::Base> kernel,
-        Image<float> &nonstationarity_image,
+        Image<float> &vst_noise_image,
         std::shared_ptr<Estimator::Base> estimator,
         filter_type filter,
         aggregator_type aggregator,

From 4416f1265fd736c1af87ae8702cb5872480cdd59 Mon Sep 17 00:00:00 2001
From: Robert Smith <robert.smith@florey.edu.au>
Date: Fri, 13 Dec 2024 12:21:26 +1100
Subject: [PATCH 33/34] mrfilter: Expose linear phase demodulation at
 command-line

Erroneously omitted from bba4fe17b4eaa152adc2cb152161f5bd73bf740c.
---
 cmd/mrfilter.cpp | 1 +
 1 file changed, 1 insertion(+)

diff --git a/cmd/mrfilter.cpp b/cmd/mrfilter.cpp
index 78d06c00ab..be2713d0c0 100644
--- a/cmd/mrfilter.cpp
+++ b/cmd/mrfilter.cpp
@@ -147,6 +147,7 @@ void usage() {
 
   OPTIONS
   + FFTAxesOption
+  + DemodulateOption
   + FFTOption
   + GradientOption
   + KSpaceOption

From 8a068a97082cc8bdfe54afd3868911fa7b17dd29 Mon Sep 17 00:00:00 2001
From: Robert Smith <robert.smith@florey.edu.au>
Date: Sun, 15 Dec 2024 17:35:39 +1100
Subject: [PATCH 34/34] dwidenoise: Add demeaning

By default, if the input image series contains a gradient table that appears to be arranged into shells, then the demeaning will be computed per shell per voxel; otherwise a single mean will be computed per voxel.
Demeaning can also be disabled using the -demean option.
The code responsible for the demeaning has been integrated with the code for phase demodulation and the variance-stabilising transform as "preconditioning".
The presence of demeaning does not affect the PCA or the noise level estimation in any way; it only influences the _reported_ signal rank.
Supersedes #2363.
---
 cmd/dwi2noise.cpp                   |  88 ++++---
 cmd/dwidenoise.cpp                  | 151 ++++++------
 core/filter/demodulate.h            |   5 +
 core/filter/kspace.h                |   2 +-
 src/denoise/demodulate.cpp          | 107 --------
 src/denoise/demodulate.h            |  46 ----
 src/denoise/denoise.h               |   1 +
 src/denoise/estimate.cpp            |  35 ---
 src/denoise/estimate.h              |   5 -
 src/denoise/estimator/estimator.cpp |   4 +-
 src/denoise/estimator/estimator.h   |   3 +-
 src/denoise/estimator/import.h      |  17 +-
 src/denoise/exports.h               |   2 -
 src/denoise/precondition.cpp        | 369 ++++++++++++++++++++++++++++
 src/denoise/precondition.h          |  95 +++++++
 src/denoise/recon.cpp               |   9 +-
 src/denoise/recon.h                 |   1 -
 17 files changed, 624 insertions(+), 316 deletions(-)
 delete mode 100644 src/denoise/demodulate.cpp
 delete mode 100644 src/denoise/demodulate.h
 create mode 100644 src/denoise/precondition.cpp
 create mode 100644 src/denoise/precondition.h

diff --git a/cmd/dwi2noise.cpp b/cmd/dwi2noise.cpp
index 8e94586bc5..5cc5508efe 100644
--- a/cmd/dwi2noise.cpp
+++ b/cmd/dwi2noise.cpp
@@ -19,12 +19,13 @@
 #include "algo/threaded_loop.h"
 #include "axes.h"
 #include "command.h"
-#include "denoise/demodulate.h"
 #include "denoise/estimate.h"
 #include "denoise/estimator/estimator.h"
 #include "denoise/exports.h"
 #include "denoise/kernel/kernel.h"
+#include "denoise/precondition.h"
 #include "denoise/subsample.h"
+#include "dwi/gradient.h"
 #include "exception.h"
 #include "filter/demodulate.h"
 
@@ -104,12 +105,10 @@ void usage() {
   + Estimator::estimator_option
   + Kernel::options
   + subsample_option
-  + demodulation_options
-  + Option("vst",
-           "apply a within-patch variance-stabilising transformation based on a pre-estimated noise level map; "
-           "note that this will be used for within-patch non-stationariy correction only, "
-           "the output noise level estimate will still be derived from the input data")
-    + Argument("image").type_image_in()
+  + precondition_options
+
+  + DWI::GradImportOptions()
+  + DWI::GradExportOptions()
 
   + OptionGroup("Options for exporting additional data regarding PCA behaviour")
   + Option("rank",
@@ -156,37 +155,58 @@ void usage() {
 // clang-format on
 
 template <typename T>
-void run(Header &data,
+void run(Image<T> &input,
          std::shared_ptr<Subsample> subsample,
          std::shared_ptr<Kernel::Base> kernel,
-         Image<float> &vst_noise_image,
          std::shared_ptr<Estimator::Base> estimator,
          Exports &exports) {
-  auto input = data.get_image<T>().with_direct_io(3);
-  Estimate<T> func(data, subsample, kernel, vst_noise_image, estimator, exports);
-  ThreadedLoop("running MP-PCA noise level estimation", data, 0, 3).run(func, input);
+  Estimate<T> func(input, subsample, kernel, estimator, exports);
+  ThreadedLoop("running MP-PCA noise level estimation", input, 0, 3).run(func, input);
 }
 
 template <typename T>
-void run(Header &data,
+void run(Header &dwi,
          const Demodulation &demodulation,
+         const demean_type demean,
+         Image<float> &vst_noise_image,
          std::shared_ptr<Subsample> subsample,
          std::shared_ptr<Kernel::Base> kernel,
-         Image<float> &vst_noise_image,
          std::shared_ptr<Estimator::Base> estimator,
          Exports &exports) {
-  if (!demodulation) {
-    run<T>(data, subsample, kernel, vst_noise_image, estimator, exports);
+  auto opt_preconditioned = get_options("preconditioned");
+  if (!demodulation && demean == demean_type::NONE && !vst_noise_image.valid()) {
+    if (!opt_preconditioned.empty()) {
+      WARN("-preconditioned option ignored: no preconditioning taking place");
+    }
+    Image<T> input = dwi.get_image<T>().with_direct_io(3);
+    run<T>(input, subsample, kernel, estimator, exports);
     return;
   }
-  auto input = data.get_image<T>();
-  auto input_demod = Image<T>::scratch(data, "Phase-demodulated version of \"" + data.name() + "\"");
-  {
-    Filter::Demodulate demodulator(input, demodulation.axes, demodulation.mode == demodulation_t::LINEAR);
-    demodulator(input, input_demod);
+  Image<T> input(dwi.get_image<T>());
+  const Precondition<T> preconditioner(input, demodulation, demean, vst_noise_image);
+  Header H_preconditioned(input);
+  Stride::set(H_preconditioned, Stride::contiguous_along_axis(3, input));
+  Image<T> input_preconditioned;
+  input_preconditioned = opt_preconditioned.empty()
+                             ? Image<T>::scratch(H_preconditioned, "Preconditioned version of \"" + input.name() + "\"")
+                             : Image<T>::create(opt_preconditioned[0][0], H_preconditioned);
+  preconditioner(input, input_preconditioned, false);
+  run(input_preconditioned, subsample, kernel, estimator, exports);
+  if (vst_noise_image.valid()) {
+    Interp::Cubic<Image<float>> vst(vst_noise_image);
+    const Transform transform(exports.noise_out);
+    for (auto l = Loop(exports.noise_out)(exports.noise_out); l; ++l) {
+      vst.scanner(transform.voxel2scanner * Eigen::Vector3d({default_type(exports.noise_out.index(0)),
+                                                             default_type(exports.noise_out.index(1)),
+                                                             default_type(exports.noise_out.index(2))}));
+      exports.noise_out.value() *= vst.value();
+    }
+  }
+  if (preconditioner.rank() == 1 && exports.rank_input.valid()) {
+    for (auto l = Loop(exports.rank_input)(exports.rank_input); l; ++l)
+      exports.rank_input.value() =
+          std::max<uint16_t>(uint16_t(exports.rank_input.value()) + uint16_t(1), uint16_t(dwi.size(3)));
   }
-  Estimate<T> func(data, subsample, kernel, vst_noise_image, estimator, exports);
-  ThreadedLoop("running MP-PCA noise level estimation", data, 0, 3).run(func, input_demod);
 }
 
 void run() {
@@ -195,7 +215,12 @@ void run() {
     throw Exception("input image must be 4-dimensional");
   bool complex = dwi.datatype().is_complex();
 
-  const Demodulation demodulation = get_demodulation(dwi);
+  const Demodulation demodulation = select_demodulation(dwi);
+  const demean_type demean = select_demean(dwi);
+  Image<float> vst_noise_image;
+  auto opt = get_options("vst");
+  if (!opt.empty())
+    vst_noise_image = Image<float>::open(opt[0][0]);
 
   auto subsample = Subsample::make(dwi);
   assert(subsample);
@@ -203,12 +228,7 @@ void run() {
   auto kernel = Kernel::make_kernel(dwi, subsample->get_factors());
   assert(kernel);
 
-  Image<float> vst_noise_image;
-  auto opt = get_options("vst");
-  if (!opt.empty())
-    vst_noise_image = Image<float>::open(opt[0][0]);
-
-  auto estimator = Estimator::make_estimator(false);
+  auto estimator = Estimator::make_estimator(vst_noise_image, false);
   assert(estimator);
 
   Exports exports(dwi, subsample->header());
@@ -233,20 +253,20 @@ void run() {
   case 0:
     assert(demodulation.axes.empty());
     INFO("select real float32 for processing");
-    run<float>(dwi, subsample, kernel, vst_noise_image, estimator, exports);
+    run<float>(dwi, demodulation, demean, vst_noise_image, subsample, kernel, estimator, exports);
     break;
   case 1:
     assert(demodulation.axes.empty());
     INFO("select real float64 for processing");
-    run<double>(dwi, subsample, kernel, vst_noise_image, estimator, exports);
+    run<double>(dwi, demodulation, demean, vst_noise_image, subsample, kernel, estimator, exports);
     break;
   case 2:
     INFO("select complex float32 for processing");
-    run<cfloat>(dwi, demodulation, subsample, kernel, vst_noise_image, estimator, exports);
+    run<cfloat>(dwi, demodulation, demean, vst_noise_image, subsample, kernel, estimator, exports);
     break;
   case 3:
     INFO("select complex float64 for processing");
-    run<cdouble>(dwi, demodulation, subsample, kernel, vst_noise_image, estimator, exports);
+    run<cdouble>(dwi, demodulation, demean, vst_noise_image, subsample, kernel, estimator, exports);
     break;
   }
 }
diff --git a/cmd/dwidenoise.cpp b/cmd/dwidenoise.cpp
index b2da3e4524..67c1ebdfcf 100644
--- a/cmd/dwidenoise.cpp
+++ b/cmd/dwidenoise.cpp
@@ -25,7 +25,6 @@
 #include <Eigen/Dense>
 #include <Eigen/Eigenvalues>
 
-#include "denoise/demodulate.h"
 #include "denoise/denoise.h"
 #include "denoise/estimator/base.h"
 #include "denoise/estimator/estimator.h"
@@ -39,6 +38,7 @@
 #include "denoise/kernel/kernel.h"
 #include "denoise/kernel/sphere_radius.h"
 #include "denoise/kernel/sphere_ratio.h"
+#include "denoise/precondition.h"
 #include "denoise/recon.h"
 #include "denoise/subsample.h"
 
@@ -146,14 +146,7 @@ void usage() {
   + Estimator::estimator_denoise_options
   + Kernel::options
   + subsample_option
-  + demodulation_options
-
-  + Option("vst",
-           "apply a within-patch variance-stabilising transformation based on a pre-estimated noise level map; "
-           "note that this will be used for within-patch non-stationariy correction only, "
-           "if noise level estimate is to be used for denoising also "
-           "it must be additionally provided via the -noise_in option")
-    + Argument("image").type_image_in()
+  + precondition_options
 
   + OptionGroup("Options that affect reconstruction of the output image series")
   + Option("filter",
@@ -201,10 +194,6 @@ void usage() {
   + Option("sum_optshrink",
            "the sum of eigenvector weights computed for the denoising patch centred at each voxel "
            "as a result of performing optimal shrinkage")
-    + Argument("image").type_image_out()
-    + Option("noise_cov",
-           "export an image of the Coefficient of Variation (CoV) of noise level within each patch "
-           "(only applicable if -nonstationarity is specified)")
     + Argument("image").type_image_out();
 
   COPYRIGHT =
@@ -242,23 +231,16 @@ void usage() {
 std::complex<double> operator/(const std::complex<double> &c, const float n) { return c / double(n); }
 
 template <typename T>
-void run(Header &data,
+void run(Image<T> &input,
          std::shared_ptr<Subsample> subsample,
          std::shared_ptr<Kernel::Base> kernel,
-         Image<float> &vst_noise_image,
          std::shared_ptr<Estimator::Base> estimator,
          filter_type filter,
          aggregator_type aggregator,
-         const std::string &output_name,
+         Image<T> &output,
          Exports &exports) {
-  auto input = data.get_image<T>().with_direct_io(3);
-  // create output
-  Header header(data);
-  header.datatype() = DataType::from<T>();
-  auto output = Image<T>::create(output_name, header);
-  // run
-  Recon<T> func(data, subsample, kernel, vst_noise_image, estimator, filter, aggregator, exports);
-  ThreadedLoop("running MP-PCA denoising", data, 0, 3).run(func, input, output);
+  Recon<T> func(input, subsample, kernel, estimator, filter, aggregator, exports);
+  ThreadedLoop("running MP-PCA denoising", input, 0, 3).run(func, input, output);
   // Rescale output if aggregation was performed
   if (aggregator == aggregator_type::EXCLUSIVE)
     return;
@@ -273,49 +255,75 @@ void run(Header &data,
 }
 
 template <typename T>
-void run(Header &data,
+void run(Header &dwi,
          const Demodulation &demodulation,
+         const demean_type demean,
+         Image<float> &vst_noise_image,
          std::shared_ptr<Subsample> subsample,
          std::shared_ptr<Kernel::Base> kernel,
-         Image<float> &vst_noise_image,
          std::shared_ptr<Estimator::Base> estimator,
          filter_type filter,
          aggregator_type aggregator,
          const std::string &output_name,
          Exports &exports) {
-  if (!demodulation) {
-    run<T>(data, subsample, kernel, vst_noise_image, estimator, filter, aggregator, output_name, exports);
+  auto opt_preconditioned = get_options("preconditioned");
+  if (!demodulation && demean == demean_type::NONE && !vst_noise_image.valid()) {
+    if (!opt_preconditioned.empty()) {
+      WARN("-preconditioned option ignored: no preconditioning taking place");
+    }
+    auto input = dwi.get_image<T>().with_direct_io(3);
+    Header H(dwi);
+    H.datatype() = DataType::from<T>();
+    auto output = Image<T>::create(output_name, H);
+    run<T>(input, subsample, kernel, estimator, filter, aggregator, output, exports);
     return;
   }
-  auto input = data.get_image<T>();
-  // generate scratch version of DWI with phase demodulation
-  Header H_scratch(data);
-  Stride::set(H_scratch, Stride::contiguous_along_axis(3));
-  H_scratch.datatype() = DataType::from<T>();
-  H_scratch.datatype().set_byte_order_native();
-  auto input_demodulated = Image<T>::scratch(H_scratch, "Phase-demodulated version of input DWI");
-  Filter::Demodulate demodulate(input, demodulation.axes, demodulation.mode == demodulation_t::LINEAR);
-  demodulate(input, input_demodulated, false);
-  input = Image<T>(); // free memory
+  auto input = dwi.get_image<T>();
+  // perform preconditioning
+  const Precondition<T> preconditioner(input, demodulation, demean, vst_noise_image);
+  Header H_preconditioned(dwi);
+  Stride::set(H_preconditioned, Stride::contiguous_along_axis(3));
+  H_preconditioned.datatype() = DataType::from<T>();
+  H_preconditioned.datatype().set_byte_order_native();
+  Image<T> input_preconditioned;
+  input_preconditioned = opt_preconditioned.empty()
+                             ? Image<T>::scratch(H_preconditioned, "Preconditioned version of \"" + dwi.name() + "\"")
+                             : Image<T>::create(opt_preconditioned[0][0], H_preconditioned);
+  preconditioner(input, input_preconditioned, false);
   // create output
-  Header header(data);
-  header.datatype() = DataType::from<T>();
-  auto output = Image<T>::create(output_name, header);
+  Header H(dwi);
+  H.datatype() = DataType::from<T>();
+  auto output = Image<T>::create(output_name, H);
   // run
-  Recon<T> func(data, subsample, kernel, vst_noise_image, estimator, filter, aggregator, exports);
-  ThreadedLoop("running MP-PCA denoising", data, 0, 3).run(func, input_demodulated, output);
-  // Re-apply phase ramps that were previously demodulated
-  demodulate(output, true);
-  // Rescale output if performing aggregation
-  if (aggregator == aggregator_type::EXCLUSIVE)
-    return;
-  for (auto l_voxel = Loop(exports.sum_aggregation)(output, exports.sum_aggregation); l_voxel; ++l_voxel) {
-    for (auto l_volume = Loop(3)(output); l_volume; ++l_volume)
-      output.value() /= float(exports.sum_aggregation.value());
+  run(input_preconditioned, subsample, kernel, estimator, filter, aggregator, output, exports);
+  // reverse effects of preconditioning
+  Image<T> output2(output);
+  preconditioner(output, output2, true);
+  // compensate for effects of preconditioning where relevant
+  if (exports.noise_out.valid() && vst_noise_image.valid()) {
+    Interp::Cubic<Image<float>> vst(vst_noise_image);
+    const Transform transform(exports.noise_out);
+    for (auto l = Loop(exports.noise_out)(exports.noise_out); l; ++l) {
+      vst.scanner(transform.voxel2scanner * Eigen::Vector3d{default_type(exports.noise_out.index(0)),
+                                                            default_type(exports.noise_out.index(1)),
+                                                            default_type(exports.noise_out.index(2))});
+      exports.noise_out.value() *= vst.value();
+    }
   }
-  if (exports.rank_output.valid()) {
-    for (auto l = Loop(exports.sum_aggregation)(exports.rank_output, exports.sum_aggregation); l; ++l)
-      exports.rank_output.value() /= exports.sum_aggregation.value();
+  if (preconditioner.rank() == 1) {
+    if (exports.rank_input.valid()) {
+      for (auto l = Loop(exports.rank_input)(exports.rank_input); l; ++l)
+        exports.rank_input.value() =
+            std::min<uint16_t>(uint16_t(exports.rank_input.value()) + uint16_t(1), uint16_t(dwi.size(3)));
+    }
+    if (exports.rank_output.valid()) {
+      for (auto l = Loop(exports.rank_output)(exports.rank_output); l; ++l)
+        exports.rank_output.value() = std::min<float>(float(exports.rank_output.value()) + 1.0f, float(dwi.size(3)));
+    }
+    if (exports.sum_optshrink.valid()) {
+      for (auto l = Loop(exports.sum_optshrink)(exports.sum_optshrink); l; ++l)
+        exports.sum_optshrink.value() = float(exports.sum_optshrink.value()) + 1.0f;
+    }
   }
 }
 
@@ -324,7 +332,12 @@ void run() {
   if (dwi.ndim() != 4 || dwi.size(3) <= 1)
     throw Exception("input image must be 4-dimensional");
 
-  const Demodulation demodulation = get_demodulation(dwi);
+  const Demodulation demodulation = select_demodulation(dwi);
+  const demean_type demean = select_demean(dwi);
+  Image<float> vst_noise_image;
+  auto opt = get_options("vst");
+  if (!opt.empty())
+    vst_noise_image = Image<float>::open(opt[0][0]);
 
   auto subsample = Subsample::make(dwi);
   assert(subsample);
@@ -332,12 +345,7 @@ void run() {
   auto kernel = Kernel::make_kernel(dwi, subsample->get_factors());
   assert(kernel);
 
-  Image<float> vst_noise_image;
-  auto opt = get_options("vst");
-  if (!opt.empty())
-    vst_noise_image = Image<float>::open(opt[0][0]);
-
-  auto estimator = Estimator::make_estimator(true);
+  auto estimator = Estimator::make_estimator(vst_noise_image, true);
   assert(estimator);
 
   filter_type filter = get_options("fixed_rank").empty() ? filter_type::OPTSHRINK : filter_type::TRUNCATE;
@@ -400,13 +408,6 @@ void run() {
     exports.set_sum_aggregation("");
   }
 
-  opt = get_options("noise_cov");
-  if (!opt.empty()) {
-    if (!vst_noise_image.valid())
-      throw Exception("-noise_variance can only be specified if -nonstationarity option is used");
-    exports.set_noise_cov(opt[0][0]);
-  }
-
   int prec = get_option_value("datatype", 0); // default: single precision
   if (dwi.datatype().is_complex())
     prec += 2; // support complex input data
@@ -416,9 +417,11 @@ void run() {
     INFO("select real float32 for processing");
     run<float>(          //
         dwi,             //
+        demodulation,    //
+        demean,          //
+        vst_noise_image, //
         subsample,       //
         kernel,          //
-        vst_noise_image, //
         estimator,       //
         filter,          //
         aggregator,      //
@@ -430,9 +433,11 @@ void run() {
     INFO("select real float64 for processing");
     run<double>(         //
         dwi,             //
+        demodulation,    //
+        demean,          //
+        vst_noise_image, //
         subsample,       //
         kernel,          //
-        vst_noise_image, //
         estimator,       //
         filter,          //
         aggregator,      //
@@ -444,9 +449,10 @@ void run() {
     run<cfloat>(         //
         dwi,             //
         demodulation,    //
+        demean,          //
+        vst_noise_image, //
         subsample,       //
         kernel,          //
-        vst_noise_image, //
         estimator,       //
         filter,          //
         aggregator,      //
@@ -458,9 +464,10 @@ void run() {
     run<cdouble>(        //
         dwi,             //
         demodulation,    //
+        demean,          //
+        vst_noise_image, //
         subsample,       //
         kernel,          //
-        vst_noise_image, //
         estimator,       //
         filter,          //
         aggregator,      //
diff --git a/core/filter/demodulate.h b/core/filter/demodulate.h
index 1fc16e55c7..aaa366d3ef 100644
--- a/core/filter/demodulate.h
+++ b/core/filter/demodulate.h
@@ -34,6 +34,9 @@ namespace MR::Filter {
 
 // From Manzano-Patron et al. 2024
 constexpr default_type default_tukey_FWHM_demodulate = 0.58;
+// TODO Ideally do more experimentation to figure out a reasonable default here
+// Too high and everything ends up in the real axis;
+//   too low and disjointed phase cound drive up signal rank
 constexpr default_type default_tukey_alpha_demodulate = 2.0 * (1.0 - default_tukey_FWHM_demodulate);
 
 /*! Estimate a linear phase ramp of a complex image and demodulate by such
@@ -245,6 +248,8 @@ class Demodulate : public Base {
     }
   }
 
+  Image<cfloat> operator()() const { return phase; }
+
 protected:
   // TODO Change to Image<float>; can produce complex value at processing time
   Image<cfloat> phase;
diff --git a/core/filter/kspace.h b/core/filter/kspace.h
index e6e99d3c45..dd5cb33d78 100644
--- a/core/filter/kspace.h
+++ b/core/filter/kspace.h
@@ -28,7 +28,7 @@
 
 namespace MR::Filter {
 
-std::vector<std::string> kspace_window_choices({"tukey"});
+const std::vector<std::string> kspace_window_choices({"tukey"});
 enum class kspace_windowfn_t { TUKEY };
 constexpr default_type default_tukey_width = 0.5;
 
diff --git a/src/denoise/demodulate.cpp b/src/denoise/demodulate.cpp
deleted file mode 100644
index f5c5a5e009..0000000000
--- a/src/denoise/demodulate.cpp
+++ /dev/null
@@ -1,107 +0,0 @@
-/* Copyright (c) 2008-2024 the MRtrix3 contributors.
- *
- * This Source Code Form is subject to the terms of the Mozilla Public
- * License, v. 2.0. If a copy of the MPL was not distributed with this
- * file, You can obtain one at http://mozilla.org/MPL/2.0/.
- *
- * Covered Software is provided under this License on an "as is"
- * basis, without warranty of any kind, either expressed, implied, or
- * statutory, including, without limitation, warranties that the
- * Covered Software is free of defects, merchantable, fit for a
- * particular purpose or non-infringing.
- * See the Mozilla Public License v. 2.0 for more details.
- *
- * For more details, see http://www.mrtrix.org/.
- */
-
-#include "denoise/demodulate.h"
-
-#include "app.h"
-#include "axes.h"
-
-using namespace MR::App;
-
-namespace MR::Denoise {
-
-const char *const demodulation_description =
-    "If the input data are of complex type, "
-    "then a smooth non-linear phase will be demodulated removed from each k-space prior to PCA. "
-    "In the absence of metadata indicating otherwise, "
-    "it is inferred that the first two axes correspond to acquired slices, "
-    "and different slices / volumes will be demodulated individually; "
-    "this behaviour can be modified using the -demod_axes option. "
-    "A strictly linear phase term can instead be regressed from each k-space, "
-    "similarly to performed in Cordero-Grande et al. 2019, "
-    "by specifying -demodulate linear.";
-
-// clang-format off
-const OptionGroup demodulation_options = OptionGroup("Options for phase demodulation of complex data")
-  + Option("demodulate",
-           "select form of phase demodulation; "
-           "options are: " + join(demodulation_choices, ",") + " "
-           "(default: nonlinear)")
-    + Argument("mode").type_choice(demodulation_choices)
-  + Option("demod_axes",
-           "comma-separated list of axis indices along which FFT can be applied for phase demodulation")
-    + Argument("axes").type_sequence_int();
-// clang-format on
-
-Demodulation get_demodulation(const Header &H) {
-  const bool complex = H.datatype().is_complex();
-  auto opt_mode = get_options("demodulate");
-  auto opt_axes = get_options("demod_axes");
-  Demodulation result;
-  if (opt_mode.empty()) {
-    if (complex) {
-      result.mode = demodulation_t::NONLINEAR;
-    } else {
-      if (!opt_axes.empty()) {
-        throw Exception("Option -demod_axes cannot be specified: "
-                        "no phase demodulation of magnitude data");
-      }
-    }
-  } else {
-    result.mode = demodulation_t(int(opt_mode[0][0]));
-    if (!complex) {
-      switch (result.mode) {
-      case demodulation_t::NONE:
-        WARN("Specifying -demodulate none is redundant: "
-             "never any phase demodulation for magnitude input data");
-        break;
-      default:
-        throw Exception("Phase modulation cannot be utilised for magnitude-only input data");
-      }
-    }
-  }
-  if (!complex)
-    return result;
-  if (opt_axes.empty()) {
-    auto slice_encoding_it = H.keyval().find("SliceEncodingDirection");
-    if (slice_encoding_it == H.keyval().end()) {
-      // TODO Ideally this would be the first two axes *on disk*,
-      //   not following transform realignment
-      INFO("No header information on slice encoding; "
-           "assuming first two axes are within-slice");
-      result.axes = {0, 1};
-    } else {
-      auto dir = Axes::id2dir(slice_encoding_it->second);
-      for (size_t axis = 0; axis != 3; ++axis) {
-        if (!dir[axis])
-          result.axes.push_back(axis);
-      }
-      INFO("For header SliceEncodingDirection=\"" + slice_encoding_it->second +
-           "\", "
-           "chose demodulation axes: " +
-           join(result.axes, ","));
-    }
-  } else {
-    result.axes = parse_ints<size_t>(opt_axes[0][0]);
-    for (auto axis : result.axes) {
-      if (axis > 2)
-        throw Exception("Phase demodulation implementation not yet robust to non-spatial axes");
-    }
-  }
-  return result;
-}
-
-} // namespace MR::Denoise
diff --git a/src/denoise/demodulate.h b/src/denoise/demodulate.h
deleted file mode 100644
index b0d9d04153..0000000000
--- a/src/denoise/demodulate.h
+++ /dev/null
@@ -1,46 +0,0 @@
-/* Copyright (c) 2008-2024 the MRtrix3 contributors.
- *
- * This Source Code Form is subject to the terms of the Mozilla Public
- * License, v. 2.0. If a copy of the MPL was not distributed with this
- * file, You can obtain one at http://mozilla.org/MPL/2.0/.
- *
- * Covered Software is provided under this License on an "as is"
- * basis, without warranty of any kind, either expressed, implied, or
- * statutory, including, without limitation, warranties that the
- * Covered Software is free of defects, merchantable, fit for a
- * particular purpose or non-infringing.
- * See the Mozilla Public License v. 2.0 for more details.
- *
- * For more details, see http://www.mrtrix.org/.
- */
-
-#pragma once
-
-#include <string>
-#include <vector>
-
-#include "app.h"
-#include "header.h"
-
-namespace MR::Denoise {
-
-extern const char *const demodulation_description;
-
-const std::vector<std::string> demodulation_choices({"none", "linear", "nonlinear"});
-enum class demodulation_t { NONE, LINEAR, NONLINEAR };
-
-extern const App::OptionGroup demodulation_options;
-
-class Demodulation {
-public:
-  Demodulation(demodulation_t mode) : mode(mode) {}
-  Demodulation() : mode(demodulation_t::NONE) {}
-  explicit operator bool() const { return mode != demodulation_t::NONE; }
-  bool operator!() const { return mode == demodulation_t::NONE; }
-  demodulation_t mode;
-  std::vector<size_t> axes;
-};
-
-Demodulation get_demodulation(const Header &);
-
-} // namespace MR::Denoise
diff --git a/src/denoise/denoise.h b/src/denoise/denoise.h
index ee17664d34..32cd596735 100644
--- a/src/denoise/denoise.h
+++ b/src/denoise/denoise.h
@@ -17,6 +17,7 @@
 #pragma once
 
 #include <Eigen/Dense>
+#include <string>
 #include <vector>
 
 #include "app.h"
diff --git a/src/denoise/estimate.cpp b/src/denoise/estimate.cpp
index 3fb8da2af7..fdc5e1617e 100644
--- a/src/denoise/estimate.cpp
+++ b/src/denoise/estimate.cpp
@@ -27,15 +27,12 @@ template <typename F>
 Estimate<F>::Estimate(const Header &header,
                       std::shared_ptr<Subsample> subsample,
                       std::shared_ptr<Kernel::Base> kernel,
-                      Image<float> &vst_noise_image,
                       std::shared_ptr<Estimator::Base> estimator,
                       Exports &exports)
     : m(header.size(3)),
       subsample(subsample),
       kernel(kernel),
       estimator(estimator),
-      transform(std::make_shared<Transform>(header)),
-      vst_noise_image(vst_noise_image),
       X(m, kernel->estimated_size()),
       XtX(std::min(m, kernel->estimated_size()), std::min(m, kernel->estimated_size())),
       eig(std::min(m, kernel->estimated_size())),
@@ -123,42 +120,10 @@ template <typename F> void Estimate<F>::operator()(Image<F> &dwi) {
       exports.patchcount.value() = exports.patchcount.value() + 1;
     }
   }
-  if (exports.noise_cov.valid()) {
-    double variance(double(0));
-    for (auto v : patch.voxels)
-      variance += Math::pow2(v.noise_level - patch.centre_noise);
-    variance /= (patch.voxels.size() - 1);
-    assign_pos_of(ss_index).to(exports.noise_cov);
-    exports.noise_cov.value() = std::sqrt(variance) / patch.centre_noise;
-  }
 }
 
 template <typename F> void Estimate<F>::load_data(Image<F> &image) {
   const Kernel::Voxel::index_type pos({image.index(0), image.index(1), image.index(2)});
-  if (vst_noise_image.valid()) {
-    assert(patch.centre_realspace.allFinite());
-    Interp::Cubic<Image<float>> interp(vst_noise_image);
-    interp.scanner(patch.centre_realspace);
-    assert(!(!interp));
-    patch.centre_noise = interp.value();
-    if (patch.centre_noise > 0.0) {
-      for (ssize_t i = 0; i != patch.voxels.size(); ++i) {
-        interp.scanner(transform->voxel2scanner * patch.voxels[i].index.cast<default_type>());
-        // TODO Trying to pull intensity information from voxels beyond the extremities of the subsampled image
-        //   may cause problems
-        assert(!(!interp));
-        const double voxel_noise = interp.value();
-        patch.voxels[i].noise_level = voxel_noise;
-        const double scaling_factor = voxel_noise > 0.0 ? (patch.centre_noise / voxel_noise) : 1.0;
-        assert(std::isfinite(scaling_factor));
-        assign_pos_of(patch.voxels[i].index, 0, 3).to(image);
-        X.col(i) = image.row(3);
-        X.col(i) *= scaling_factor;
-      }
-      assign_pos_of(pos, 0, 3).to(image);
-      return;
-    }
-  }
   for (ssize_t i = 0; i != patch.voxels.size(); ++i) {
     assign_pos_of(patch.voxels[i].index, 0, 3).to(image);
     X.col(i) = image.row(3);
diff --git a/src/denoise/estimate.h b/src/denoise/estimate.h
index 622acdd7f7..7aa9f46589 100644
--- a/src/denoise/estimate.h
+++ b/src/denoise/estimate.h
@@ -43,7 +43,6 @@ template <typename F> class Estimate {
   Estimate(const Header &header,
            std::shared_ptr<Subsample> subsample,
            std::shared_ptr<Kernel::Base> kernel,
-           Image<float> &vst_noise_image,
            std::shared_ptr<Estimator::Base> estimator,
            Exports &exports);
 
@@ -57,12 +56,8 @@ template <typename F> class Estimate {
   std::shared_ptr<Kernel::Base> kernel;
   std::shared_ptr<Estimator::Base> estimator;
 
-  // Necessary for transform from input voxel locations to nonstationarity image
-  std::shared_ptr<Transform> transform;
-
   // Reusable memory
   Kernel::Data patch;
-  Image<float> vst_noise_image;
   MatrixType X;
   MatrixType XtX;
   Eigen::SelfAdjointEigenSolver<MatrixType> eig;
diff --git a/src/denoise/estimator/estimator.cpp b/src/denoise/estimator/estimator.cpp
index 41ffee1406..bf66aaffa5 100644
--- a/src/denoise/estimator/estimator.cpp
+++ b/src/denoise/estimator/estimator.cpp
@@ -53,7 +53,7 @@ const OptionGroup estimator_denoise_options =
              "set a fixed input signal rank rather than estimating the noise level from the data")
       + Argument("value").type_integer(1);
 
-std::shared_ptr<Base> make_estimator(const bool permit_bypass) {
+std::shared_ptr<Base> make_estimator(Image<float> &vst_noise_in, const bool permit_bypass) {
   auto opt = get_options("estimator");
   if (permit_bypass) {
     auto noise_in = get_options("noise_in");
@@ -63,7 +63,7 @@ std::shared_ptr<Base> make_estimator(const bool permit_bypass) {
         throw Exception("Cannot both provide an input noise level image and specify a noise level estimator");
       if (!fixed_rank.empty())
         throw Exception("Cannot both provide an input noise level image and request a fixed signal rank");
-      return std::make_shared<Import>(noise_in[0][0]);
+      return std::make_shared<Import>(noise_in[0][0], vst_noise_in);
     }
     if (!fixed_rank.empty()) {
       if (!opt.empty())
diff --git a/src/denoise/estimator/estimator.h b/src/denoise/estimator/estimator.h
index 322d65ca96..1384b157ee 100644
--- a/src/denoise/estimator/estimator.h
+++ b/src/denoise/estimator/estimator.h
@@ -21,6 +21,7 @@
 #include <vector>
 
 #include "app.h"
+#include "image.h"
 
 namespace MR::Denoise::Estimator {
 
@@ -30,6 +31,6 @@ extern const App::Option estimator_option;
 extern const App::OptionGroup estimator_denoise_options;
 const std::vector<std::string> estimators = {"exp1", "exp2", "med", "mrm2022"};
 enum class estimator_type { EXP1, EXP2, MED, MRM2022 };
-std::shared_ptr<Base> make_estimator(const bool permit_bypass);
+std::shared_ptr<Base> make_estimator(Image<float> &vst_noise_in, const bool permit_bypass);
 
 } // namespace MR::Denoise::Estimator
diff --git a/src/denoise/estimator/import.h b/src/denoise/estimator/import.h
index 6705c11155..fde0801bba 100644
--- a/src/denoise/estimator/import.h
+++ b/src/denoise/estimator/import.h
@@ -28,7 +28,9 @@ namespace MR::Denoise::Estimator {
 
 class Import : public Base {
 public:
-  Import(const std::string &path) : noise_image(Image<float>::open(path)) {}
+  Import(const std::string &path, Image<float> &vst_noise_in) //
+      : noise_image(Image<float>::open(path)),                //
+        vst_noise_image(vst_noise_in) {}                      //
   Result operator()(const eigenvalues_type &s,                //
                     const ssize_t m,                          //
                     const ssize_t n,                          //
@@ -46,7 +48,17 @@ class Import : public Base {
       //   where the patch centre is too close to the image edge for cubic interpolation
       if (!interp.scanner(pos))
         return result;
-      result.sigma2 = Math::pow2(interp.value());
+      // If the data have been preconditioned at input based on a pre-estimated noise level,
+      //   then we need to rescale the threshold that we load from this image
+      //   based on knowledge of that rescaling
+      if (vst_noise_image.valid()) {
+        Interp::Cubic<Image<float>> vst_interp(vst_noise_image);
+        if (!vst_interp.scanner(pos))
+          return result;
+        result.sigma2 = Math::pow2(interp.value() / vst_interp.value());
+      } else {
+        result.sigma2 = Math::pow2(interp.value());
+      }
     }
     // From this noise level,
     //   estimate the upper bound of the MP distribution and rank of signal
@@ -69,6 +81,7 @@ class Import : public Base {
 
 private:
   Image<float> noise_image;
+  Image<float> vst_noise_image;
 };
 
 } // namespace MR::Denoise::Estimator
diff --git a/src/denoise/exports.h b/src/denoise/exports.h
index 0cd1e8a715..d93008025c 100644
--- a/src/denoise/exports.h
+++ b/src/denoise/exports.h
@@ -59,7 +59,6 @@ class Exports {
     else
       sum_aggregation = Image<float>::create(path, H_in);
   }
-  void set_noise_cov(const std::string &path) { noise_cov = Image<float>::create(path, H_ss); }
 
   Image<float> noise_out;
   Image<uint16_t> rank_input;
@@ -69,7 +68,6 @@ class Exports {
   Image<uint16_t> voxelcount;
   Image<uint16_t> patchcount;
   Image<float> sum_aggregation;
-  Image<float> noise_cov;
 
 protected:
   Header H_in;
diff --git a/src/denoise/precondition.cpp b/src/denoise/precondition.cpp
new file mode 100644
index 0000000000..6c6c9afada
--- /dev/null
+++ b/src/denoise/precondition.cpp
@@ -0,0 +1,369 @@
+/* Copyright (c) 2008-2024 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#include "denoise/precondition.h"
+
+#include <limits>
+
+#include "algo/copy.h"
+#include "app.h"
+#include "axes.h"
+#include "dwi/gradient.h"
+#include "dwi/shells.h"
+#include "transform.h"
+
+using namespace MR::App;
+
+namespace MR::Denoise {
+
+const char *const demodulation_description =
+    "If the input data are of complex type, "
+    "then a smooth non-linear phase will be demodulated removed from each k-space prior to PCA. "
+    "In the absence of metadata indicating otherwise, "
+    "it is inferred that the first two axes correspond to acquired slices, "
+    "and different slices / volumes will be demodulated individually; "
+    "this behaviour can be modified using the -demod_axes option. "
+    "A strictly linear phase term can instead be regressed from each k-space, "
+    "similarly to performed in Cordero-Grande et al. 2019, "
+    "by specifying -demodulate linear.";
+
+// clang-format off
+const OptionGroup precondition_options = OptionGroup("Options for preconditioning data prior to PCA")
+  + Option("demodulate",
+           "select form of phase demodulation; "
+           "options are: " + join(demodulation_choices, ",") + " "
+           "(default: nonlinear)")
+    + Argument("mode").type_choice(demodulation_choices)
+  + Option("demod_axes",
+           "comma-separated list of axis indices along which FFT can be applied for phase demodulation")
+    + Argument("axes").type_sequence_int()
+  + Option("demean",
+           "select method of demeaning prior to PCA; "
+           "options are: " + join(demean_choices, ",") + " "
+           "(default: 'shells' if DWI gradient table available, 'all' otherwise)")
+    + Argument("mode").type_choice(demean_choices)
+  + Option("vst",
+           "apply a within-patch variance-stabilising transformation based on a pre-estimated noise level map")
+    + Argument("image").type_image_in()
+  + Option("preconditioned",
+           "export the preconditioned version of the input image that is the input to PCA")
+    + Argument("image").type_image_out();
+// clang-format on
+
+Demodulation select_demodulation(const Header &H) {
+  const bool complex = H.datatype().is_complex();
+  auto opt_mode = get_options("demodulate");
+  auto opt_axes = get_options("demod_axes");
+  Demodulation result;
+  if (opt_mode.empty()) {
+    if (complex) {
+      result.mode = demodulation_t::NONLINEAR;
+    } else {
+      if (!opt_axes.empty()) {
+        throw Exception("Option -demod_axes cannot be specified: "
+                        "no phase demodulation of magnitude data");
+      }
+    }
+  } else {
+    result.mode = demodulation_t(int(opt_mode[0][0]));
+    if (!complex) {
+      switch (result.mode) {
+      case demodulation_t::NONE:
+        WARN("Specifying -demodulate none is redundant: "
+             "never any phase demodulation for magnitude input data");
+        break;
+      default:
+        throw Exception("Phase modulation cannot be utilised for magnitude-only input data");
+      }
+    }
+  }
+  if (!complex)
+    return result;
+  if (opt_axes.empty()) {
+    auto slice_encoding_it = H.keyval().find("SliceEncodingDirection");
+    if (slice_encoding_it == H.keyval().end()) {
+      // TODO Ideally this would be the first two axes *on disk*,
+      //   not following transform realignment
+      INFO("No header information on slice encoding; "
+           "assuming first two axes are within-slice");
+      result.axes = {0, 1};
+    } else {
+      auto dir = Axes::id2dir(slice_encoding_it->second);
+      for (size_t axis = 0; axis != 3; ++axis) {
+        if (!dir[axis])
+          result.axes.push_back(axis);
+      }
+      INFO("For header SliceEncodingDirection=\"" + slice_encoding_it->second +
+           "\", "
+           "chose demodulation axes: " +
+           join(result.axes, ","));
+    }
+  } else {
+    result.axes = parse_ints<size_t>(opt_axes[0][0]);
+    for (auto axis : result.axes) {
+      if (axis > 2)
+        throw Exception("Phase demodulation implementation not yet robust to non-spatial axes");
+    }
+  }
+  return result;
+}
+
+demean_type select_demean(const Header &H) {
+  auto opt = get_options("demean");
+  if (opt.empty()) {
+    try {
+      auto grad = DWI::get_DW_scheme(H);
+      auto shells = DWI::Shells(grad);
+      INFO("Choosing to demean per b-value shell based on input gradient table");
+      return demean_type::SHELLS;
+    } catch (Exception &) {
+      INFO("Choosing to demean across all volumes based on absent / non-shelled gradient table");
+      return demean_type::ALL;
+    }
+  }
+  return demean_type(int(opt[0][0]));
+}
+
+template <typename T>
+Precondition<T>::Precondition(Image<T> &image,
+                              const Demodulation &demodulation,
+                              const demean_type demean,
+                              Image<float> &vst_image)
+    : H(image),              //
+      vst_image(vst_image) { //
+
+  // Step 1: Phase demodulation
+  Image<T> dephased;
+  if (demodulation.mode == demodulation_t::NONE) {
+    dephased = image;
+  } else {
+    typename DemodulatorSelector<T>::type demodulator(image,                                        //
+                                                      demodulation.axes,                            //
+                                                      demodulation.mode == demodulation_t::LINEAR); //
+    phase_image = demodulator();
+    // Only actually perform the dephasing of the input image
+    //   if that result needs to be utilised in calculation of the mean
+    if (demean != demean_type::NONE) {
+      dephased = Image<T>::scratch(H, "Scratch dephased version of \"" + image.name() + "\" for mean calculation");
+      demodulator(image, dephased, false);
+    }
+  }
+
+  // Step 2: Demeaning
+  Header H_mean(H);
+  switch (demean) {
+  case demean_type::NONE:
+    break;
+  case demean_type::SHELLS: {
+    Eigen::Matrix<default_type, Eigen::Dynamic, Eigen::Dynamic> grad;
+    try {
+      grad = DWI::get_DW_scheme(H_mean);
+    } catch (Exception &e) {
+      throw Exception(e, "Cannot demean by shells as unable to obtain valid gradient table");
+    }
+    try {
+      DWI::Shells shells(grad);
+      vol2shellidx.resize(image.size(3), -1);
+      for (ssize_t shell_idx = 0; shell_idx != shells.count(); ++shell_idx) {
+        for (auto v : shells[shell_idx].get_volumes())
+          vol2shellidx[v] = shell_idx;
+      }
+      assert(*std::min_element(vol2shellidx.begin(), vol2shellidx.end()) == 0);
+      H_mean.size(3) = shells.count();
+      DWI::stash_DW_scheme(H_mean, grad);
+      mean_image = Image<T>::scratch(H_mean, "Scratch image for per-shell mean intensity");
+      for (auto l_voxel = Loop("Computing mean intensities across shells", H_mean, 0, 3)(dephased, mean_image); //
+           l_voxel;                                                                                             //
+           ++l_voxel) {                                                                                         //
+        for (ssize_t volume_idx = 0; volume_idx != image.size(3); ++volume_idx) {
+          dephased.index(3) = volume_idx;
+          mean_image.index(3) = vol2shellidx[volume_idx];
+          mean_image.value() += dephased.value();
+        }
+        for (ssize_t shell_idx = 0; shell_idx != shells.count(); ++shell_idx) {
+          mean_image.index(3) = shell_idx;
+          mean_image.value() /= T(shells[shell_idx].count());
+        }
+      }
+    } catch (Exception &e) {
+      throw Exception(e, "Cannot demean by shells as unable to establish b-value shell structure");
+    }
+  } break;
+  case demean_type::ALL: {
+    H_mean.ndim() = 3;
+    DWI::clear_DW_scheme(H_mean);
+    mean_image = Image<T>::scratch(H_mean, "Scratch image for mean intensity across all volumes");
+    for (auto l_voxel = Loop("Computing mean intensity across all volumes", H_mean)(dephased, mean_image); //
+         l_voxel;                                                                                          //
+         ++l_voxel) {                                                                                      //
+      T mean(T(0));
+      for (auto l_volume = Loop(3)(dephased); l_volume; ++l_volume)
+        mean += T(dephased.value());
+      mean_image.value() = mean / T(image.size(3));
+    }
+  } break;
+  }
+
+  // Step 3: Variance-stabilising transform
+  // Image<float> vst is already set within constructor definition;
+  //   nothing to do here
+}
+
+namespace {
+// Private functions to prevent compiler attempting to create complex functions for real types
+template <typename T>
+typename std::enable_if<std::is_same<T, cfloat>::value, T>::type demodulate(const cfloat in, const cfloat phase) {
+  return in * std::conj(phase);
+}
+template <typename T>
+typename std::enable_if<std::is_same<T, cdouble>::value, T>::type demodulate(const cdouble in, const cfloat phase) {
+  return in * std::conj(cdouble(phase));
+}
+template <typename T>
+typename std::enable_if<!is_complex<T>::value, T>::type demodulate(const T in, const cfloat phase) {
+  assert(false);
+  return in;
+}
+template <typename T>
+typename std::enable_if<std::is_same<T, cfloat>::value, T>::type modulate(const cfloat in, const cfloat phase) {
+  return in * phase;
+}
+template <typename T>
+typename std::enable_if<std::is_same<T, cdouble>::value, T>::type modulate(const cdouble in, const cfloat phase) {
+  return in * cdouble(phase);
+}
+template <typename T> typename std::enable_if<!is_complex<T>::value, T>::type modulate(const T in, const cfloat phase) {
+  assert(false);
+  return in;
+}
+} // namespace
+
+template <typename T> void Precondition<T>::operator()(Image<T> input, Image<T> output, const bool inverse) const {
+
+  // For thread-safety / const-ness
+  const Transform transform(input);
+  Image<cfloat> phase(phase_image);
+  Image<T> mean(mean_image);
+  std::unique_ptr<Interp::Cubic<Image<float>>> vst;
+  if (vst_image.valid())
+    vst.reset(new Interp::Cubic<Image<float>>(vst_image));
+
+  Eigen::Array<T, Eigen::Dynamic, 1> data(input.size(3));
+  if (inverse) {
+    for (auto l_voxel = Loop("Reversing data preconditioning", H, 0, 3)(input, output); l_voxel; ++l_voxel) {
+
+      // Step 3: Reverse variance-stabilising transform
+      if (vst) {
+        vst->scanner(transform.voxel2scanner *                         //
+                     Eigen::Vector3d({default_type(input.index(0)),    //
+                                      default_type(input.index(1)),    //
+                                      default_type(input.index(2))})); //
+        const T multiplier = T(vst->value());
+        for (ssize_t v = 0; v != input.size(3); ++v) {
+          input.index(3) = v;
+          data[v] = T(input.value()) * multiplier;
+        }
+      } else {
+        for (ssize_t v = 0; v != input.size(3); ++v) {
+          input.index(3) = v;
+          data[v] = input.value();
+        }
+      }
+
+      // Step 2: Reverse demeaning
+      if (mean.valid()) {
+        assign_pos_of(input, 0, 3).to(mean);
+        if (mean.ndim() == 3) {
+          const T mean_value = mean.value();
+          data += mean_value;
+        } else {
+          for (ssize_t v = 0; v != input.size(3); ++v) {
+            mean.index(3) = vol2shellidx[v];
+            data[v] += T(mean.value());
+          }
+        }
+      }
+
+      // Step 1: Reverse phase demodulation
+      if (phase.valid()) {
+        assign_pos_of(input, 0, 3).to(phase);
+        for (ssize_t v = 0; v != input.size(3); ++v) {
+          phase.index(3) = v;
+          data[v] = modulate<T>(data[v], phase.value());
+        }
+      }
+
+      // Write to output
+      for (ssize_t v = 0; v != input.size(3); ++v) {
+        output.index(3) = v;
+        output.value() = data[v];
+      }
+    }
+    return;
+  }
+
+  // Applying forward preconditioning
+  for (auto l_voxel = Loop("Applying data preconditioning", H, 0, 3)(input, output); l_voxel; ++l_voxel) {
+
+    // Step 1: Phase demodulation
+    if (phase.valid()) {
+      assign_pos_of(input, 0, 3).to(phase);
+      for (ssize_t v = 0; v != input.size(3); ++v) {
+        input.index(3) = v;
+        phase.index(3) = v;
+        data[v] = demodulate<T>(input.value(), phase.value());
+      }
+    } else {
+      for (ssize_t v = 0; v != input.size(3); ++v) {
+        input.index(3) = v;
+        data[v] = input.value();
+      }
+    }
+
+    // Step 2: Demeaning
+    if (mean.valid()) {
+      assign_pos_of(input, 0, 3).to(mean);
+      if (mean.ndim() == 3) {
+        const T mean_value = mean.value();
+        for (ssize_t v = 0; v != input.size(3); ++v)
+          data[v] -= mean_value;
+      } else {
+        for (ssize_t v = 0; v != input.size(3); ++v) {
+          mean.index(3) = vol2shellidx[v];
+          data[v] -= T(mean.value());
+        }
+      }
+    }
+
+    // Step 3: Variance-stabilising transform
+    if (vst) {
+      vst->scanner(transform.voxel2scanner                             //
+                   * Eigen::Vector3d({default_type(input.index(0)),    //
+                                      default_type(input.index(1)),    //
+                                      default_type(input.index(2))})); //
+      const default_type multiplier = 1.0 / vst->value();
+      data *= multiplier;
+    }
+
+    // Write to output
+    for (ssize_t v = 0; v != input.size(3); ++v) {
+      output.index(3) = v;
+      output.value() = data[v];
+    }
+  }
+}
+
+} // namespace MR::Denoise
diff --git a/src/denoise/precondition.h b/src/denoise/precondition.h
new file mode 100644
index 0000000000..7a1ff9c74c
--- /dev/null
+++ b/src/denoise/precondition.h
@@ -0,0 +1,95 @@
+/* Copyright (c) 2008-2024 the MRtrix3 contributors.
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/.
+ *
+ * Covered Software is provided under this License on an "as is"
+ * basis, without warranty of any kind, either expressed, implied, or
+ * statutory, including, without limitation, warranties that the
+ * Covered Software is free of defects, merchantable, fit for a
+ * particular purpose or non-infringing.
+ * See the Mozilla Public License v. 2.0 for more details.
+ *
+ * For more details, see http://www.mrtrix.org/.
+ */
+
+#pragma once
+
+#include <string>
+#include <vector>
+
+#include "app.h"
+#include "denoise/kernel/voxel.h"
+#include "filter/demodulate.h"
+#include "header.h"
+#include "image.h"
+#include "interp/cubic.h"
+#include "types.h"
+
+namespace MR::Denoise {
+
+extern const char *const demodulation_description;
+
+const std::vector<std::string> demodulation_choices({"none", "linear", "nonlinear"});
+enum class demodulation_t { NONE, LINEAR, NONLINEAR };
+
+const std::vector<std::string> demean_choices = {"none", "shells", "all"};
+enum class demean_type { NONE, SHELLS, ALL };
+
+extern const App::OptionGroup precondition_options;
+
+class Demodulation {
+public:
+  Demodulation(demodulation_t mode) : mode(mode) {}
+  Demodulation() : mode(demodulation_t::NONE) {}
+  explicit operator bool() const { return mode != demodulation_t::NONE; }
+  bool operator!() const { return mode == demodulation_t::NONE; }
+  demodulation_t mode;
+  std::vector<size_t> axes;
+};
+Demodulation select_demodulation(const Header &);
+
+demean_type select_demean(const Header &);
+
+// Need to SFINAE define the demodulator type,
+//   so that it does not attempt to compile the demodulation filter for non-complex types
+class DummyDemodulator {
+public:
+  template <class ImageType> DummyDemodulator(ImageType &, const std::vector<size_t> &, const bool) {}
+  template <class InputImageType, class OutputImageType>
+  void operator()(InputImageType &, OutputImageType &, const bool) {
+    assert(false);
+  }
+  Image<cfloat> operator()() { return Image<cfloat>(); }
+};
+template <typename T> struct DemodulatorSelector {
+  using type = DummyDemodulator;
+};
+template <typename T> struct DemodulatorSelector<std::complex<T>> {
+  using type = Filter::Demodulate;
+};
+
+template <typename T> class Precondition {
+public:
+  Precondition(Image<T> &image, const Demodulation &demodulation, const demean_type demean, Image<float> &vst);
+  Precondition(Precondition &) = default;
+  void operator()(Image<T> input, Image<T> output, const bool inverse = false) const;
+  ssize_t rank() const { return phase_image.valid() || mean_image.valid() ? 1 : 0; }
+
+private:
+  const Header H;
+  // First step: Phase demodulation
+  Image<cfloat> phase_image;
+  // Second step: Demeaning
+  std::vector<ssize_t> vol2shellidx;
+  Image<T> mean_image;
+  // Third step: Variance-stabilising transform
+  Image<float> vst_image;
+};
+template class Precondition<float>;
+template class Precondition<double>;
+template class Precondition<cfloat>;
+template class Precondition<cdouble>;
+
+} // namespace MR::Denoise
diff --git a/src/denoise/recon.cpp b/src/denoise/recon.cpp
index 9cfc177667..2573e83153 100644
--- a/src/denoise/recon.cpp
+++ b/src/denoise/recon.cpp
@@ -24,12 +24,11 @@ template <typename F>
 Recon<F>::Recon(const Header &header,
                 std::shared_ptr<Subsample> subsample,
                 std::shared_ptr<Kernel::Base> kernel,
-                Image<float> &vst_noise_image,
                 std::shared_ptr<Estimator::Base> estimator,
                 filter_type filter,
                 aggregator_type aggregator,
                 Exports &exports)
-    : Estimate<F>(header, subsample, kernel, vst_noise_image, estimator, exports),
+    : Estimate<F>(header, subsample, kernel, estimator, exports),
       filter(filter),
       aggregator(aggregator),
       // FWHM = 2 x cube root of spacings between kernels
@@ -75,12 +74,6 @@ template <typename F> void Recon<F>::operator()(Image<F> &dwi, Image<F> &out) {
       const double transition = 1.0 + std::sqrt(beta);
       for (ssize_t i = 0; i != r; ++i) {
         const double lam = std::max(Estimate<F>::s[i], 0.0) / q;
-        // TODO Should this be based on the noise level,
-        //   or on the estimated upper bound of the MP distribution?
-        // If based on upper bound,
-        //   there will be an issue with importing this information from a pre-estimated noise map
-        // TODO Unexpected absence of sqrt() here
-        // const double y = lam / std::sqrt(Estimate<F>::threshold.sigma2);
         const double y = lam / Estimate<F>::threshold.sigma2;
         double nu = 0.0;
         if (y > transition) {
diff --git a/src/denoise/recon.h b/src/denoise/recon.h
index 3d96dedc3d..6db9e40d59 100644
--- a/src/denoise/recon.h
+++ b/src/denoise/recon.h
@@ -37,7 +37,6 @@ template <typename F> class Recon : public Estimate<F> {
   Recon(const Header &header,
         std::shared_ptr<Subsample> subsample,
         std::shared_ptr<Kernel::Base> kernel,
-        Image<float> &vst_noise_image,
         std::shared_ptr<Estimator::Base> estimator,
         filter_type filter,
         aggregator_type aggregator,