dwidenoise: Add option -demean

cmd/dwidenoise.cpp

@@ -47,7 +47,11 @@ 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.";
+      "reduce such non-Gaussian biases, and the command now supports complex input data."
+
+    + "If using this command to denoise non-DWI data, e.g. fMRI, it may be beneficial to use the -demean "
+      "option, which mitigates floating-point precision issues when the variance between image volumes is "
+      "small relative to the mean across volumes.";
 
   AUTHOR = "Daan Christiaens ([email protected]) & "
            "Jelle Veraart ([email protected]) & "
@@ -97,7 +101,9 @@ void usage ()
     + 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);
+    +   Argument ("Exp1/Exp2").type_choice(estimators)
+
+    + Option ("demean", "Pre-allocate the mean signal in each voxel as a signal component");
 
 
   COPYRIGHT = "Copyright (c) 2016 New York University, University of Antwerp, and the MRtrix3 contributors \n \n"
@@ -130,18 +136,22 @@ class DenoisingFunctor {
 public:
 
   using MatrixType = Eigen::Matrix<F, Eigen::Dynamic, Eigen::Dynamic>;
+  using VectorType = Eigen::Matrix<F, Eigen::Dynamic, 1>;
   using SValsType = Eigen::VectorXd;
 
   DenoisingFunctor (int ndwi, const vector<uint32_t>& extent,
                     Image<bool>& mask, Image<real_type>& noise,
-                    Image<uint16_t>& rank, bool exp1)
+                    Image<uint16_t>& rank, bool exp1, bool demean)
     : extent {{extent[0]/2, extent[1]/2, extent[2]/2}},
       m (ndwi),
       n (extent[0]*extent[1]*extent[2]),
       r (std::min(m,n)),
       q (std::max(m,n)),
-      exp1(exp1),
+      q_offset (demean ? q-1 : q),
+      exp1 (exp1),
+      demean (demean),
       X (m,n),
+      Xm (VectorType::Zero (r)),
       XtX (r, r),
       eig (r),
       s (r),
@@ -165,22 +175,31 @@ class DenoisingFunctor {
     load_data (dwi);
 
     // Compute Eigendecomposition:
-    if (m <= n)
+    if (m <= n) {
+      if (demean) {
+        Xm = X.rowwise().mean();
+        X.colwise() -= Xm;
+      }
       XtX.template triangularView<Eigen::Lower>() = X * X.adjoint();
-    else
+    } else {
+      if (demean) {
+        Xm = X.colwise().mean();
+        X.rowwise() -= Xm.transpose();
+      }
       XtX.template triangularView<Eigen::Lower>() = X.adjoint() * X;
+    }
     eig.compute (XtX);
     // eigenvalues sorted in increasing order:
     s = eig.eigenvalues().template cast<double>();
 
     // Marchenko-Pastur optimal threshold
-    const double lam_r = std::max(s[0], 0.0) / q;
+    const double lam_r = std::max(s[0], 0.0) / q_offset;
     double clam = 0.0;
     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)
-      double lam = std::max(s[p], 0.0) / q;
+      double lam = std::max(s[p], 0.0) / q_offset;
       clam += lam;
       double gam = double(p+1) / (exp1 ? q : q-(r-p-1));
       double sigsq1 = clam / double(p+1);
@@ -202,6 +221,9 @@ class DenoisingFunctor {
         X.col (n/2) = X * ( eig.eigenvectors() * ( s.cast<F>().asDiagonal() * eig.eigenvectors().adjoint().col(n/2) ));
     }
 
+    if (demean)
+      X.col (n/2).array() += Xm[n/2];
+
     // Store output
     assign_pos_of(dwi).to(out);
     out.row(3) = X.col(n/2);
@@ -214,16 +236,17 @@ class DenoisingFunctor {
     // store rank map if requested:
     if (rankmap.valid()) {
       assign_pos_of(dwi, 0, 3).to(rankmap);
-      rankmap.value() = uint16_t (r - cutoff_p);
+      rankmap.value() = uint16_t (r - cutoff_p + (demean ? 1 : 0));
     }
 
   }
 
 private:
   const std::array<ssize_t, 3> extent;
-  const ssize_t m, n, r, q;
-  const bool exp1;
+  const ssize_t m, n, r, q, q_offset;
+  const bool exp1, demean;
   MatrixType X;
+  VectorType Xm;
   MatrixType XtX;
   Eigen::SelfAdjointEigenSolver<MatrixType> eig;
   SValsType s;
@@ -268,15 +291,15 @@ class DenoisingFunctor {
 
 template <typename T>
 void process_image (Header& data, Image<bool>& mask, Image<real_type>& noise, Image<uint16_t>& rank,
-                    const std::string& output_name, const vector<uint32_t>& extent, bool exp1)
+                    const std::string& output_name, const vector<uint32_t>& extent, bool exp1, bool demean)
   {
     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), extent, mask, noise, rank, exp1);
+    DenoisingFunctor<T> func (data.size(3), extent, mask, noise, rank, exp1, demean);
     ThreadedLoop ("running MP-PCA denoising", data, 0, 3).run (func, input, output);
   }
 
@@ -346,24 +369,26 @@ void run ()
     rank = Image<uint16_t>::create (opt[0][0], header);
   }
 
+  const bool demean = get_options ("demean").size();
+
   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");
-      process_image<float>(dwi, mask, noise, rank, argument[1], extent, exp1);
+      process_image<float>(dwi, mask, noise, rank, argument[1], extent, exp1, demean);
       break;
     case 1:
       INFO("select real float64 for processing");
-      process_image<double>(dwi, mask, noise, rank, argument[1], extent, exp1);
+      process_image<double>(dwi, mask, noise, rank, argument[1], extent, exp1, demean);
       break;
     case 2:
       INFO("select complex float32 for processing");
-      process_image<cfloat>(dwi, mask, noise, rank, argument[1], extent, exp1);
+      process_image<cfloat>(dwi, mask, noise, rank, argument[1], extent, exp1, demean);
       break;
     case 3:
       INFO("select complex float64 for processing");
-      process_image<cdouble>(dwi, mask, noise, rank, argument[1], extent, exp1);
+      process_image<cdouble>(dwi, mask, noise, rank, argument[1], extent, exp1, demean);
       break;
   }
 

docs/reference/commands/dwidenoise.rst

@@ -27,6 +27,8 @@ 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 using this command to denoise non-DWI data, e.g. fMRI, it may be beneficial to use the -demean option, which mitigates floating-point precision issues when the variance between image volumes is small relative to the mean across volumes.
+
 Options
 -------
 
@@ -36,12 +38,16 @@ Options
 
 -  **-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.
 
+-  **-rank cutoff** The selected signal rank of the output denoised image.
+
 -  **-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 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).
 
+-  **-demean** Pre-allocate the mean signal in each voxel as a signal component
+
 Standard options
 ^^^^^^^^^^^^^^^^