iFOD2.h

Go to the documentation of this file.

/* Copyright (c) 2008-2022 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/.
 */
 
#ifndef __dwi_tractography_algorithms_iFOD2_h__
#define __dwi_tractography_algorithms_iFOD2_h__
 
#include <algorithm>
 
#include "types.h"
#include "math/SH.h"
#include "dwi/tractography/properties.h"
#include "dwi/tractography/tracking/method.h"
#include "dwi/tractography/tracking/shared.h"
#include "dwi/tractography/tracking/tractography.h"
#include "dwi/tractography/tracking/types.h"
#include "dwi/tractography/algorithms/calibrator.h"
 
 
namespace MR
{
  namespace DWI
  {
    namespace Tractography
    {
      namespace Algorithms
      {
 
        extern const App::OptionGroup iFOD2Options;
        void load_iFOD2_options (Tractography::Properties&);
 
        using namespace MR::DWI::Tractography::Tracking;
 
        class iFOD2 : public MethodBase { MEMALIGN(iFOD2)
          public:
 
            class Shared : public SharedBase { MEMALIGN(Shared)
              public:
                Shared (const std::string& diff_path, DWI::Tractography::Properties& property_set) :
                    SharedBase (diff_path, property_set),
                    lmax (Math::SH::LforN (source.size(3))),
                    num_samples (Defaults::ifod2_nsamples),
                    max_trials (Defaults::max_trials_per_step),
                    sin_max_angle_ho (NaN),
                    mean_samples (0.0),
                    mean_truncations (0.0),
                    max_max_truncation (0.0),
                    num_proc (0)
                {
                  try {
                    Math::SH::check (source);
                  } catch (Exception& e) {
                    e.display();
                    throw Exception ("Algorithm iFOD2 expects as input a spherical harmonic (SH) image");
                  }
 
                  if (rk4)
                    throw Exception ("4th-order Runge-Kutta integration not valid for iFOD2 algorithm");
 
                  set_step_and_angle (Defaults::stepsize_voxels_ifod2, Defaults::angle_ifod2, true);
                  sin_max_angle_ho = std::sin (max_angle_ho);
                  set_cutoff (Defaults::cutoff_fod * (is_act() ? Defaults::cutoff_act_multiplier : 1.0));
 
                  properties["method"] = "iFOD2";
                  properties.set (lmax, "lmax");
                  properties.set (num_samples, "samples_per_step");
                  properties.set (max_trials, "max_trials");
                  fod_power = 1.0/num_samples;
                  properties.set (fod_power, "fod_power");
                  bool precomputed = true;
                  properties.set (precomputed, "sh_precomputed");
                  if (precomputed)
                    precomputer.init (lmax);
 
                  // num_samples is number of samples excluding first point
                  --num_samples;
                  INFO ("iFOD2 generating " + str(num_samples) + " vertices per " + str (step_size) + " mm step");
 
                  // iFOD2 by default downsamples after track propagation back to the desired 'step size'
                  //   i.e. the sub-step detail is removed from the output
                  size_t downsample_factor = num_samples;
                  properties.set (downsample_factor, "downsample_factor");
                  downsampler.set_ratio (downsample_factor);
 
                  // For iFOD2, "step_size" represents the length of the chord represented
                  //   using "num_samples" vertices rather than just one; the following two
                  //   variables need to be calculated accordingly:
                  //   - The arc angle subtended by two sequential vertices on a circle of minimal radius
                  //     (prior to downsampling)
                  const float angle_minradius_preds = 2.0 * std::asin (step_size / (2.0 * min_radius)) / float(num_samples);
                  //   - The maximal possible distance between vertices after downsampling
                  const float max_step_postds = downsample_factor * step_size / float(num_samples);
                  set_num_points (angle_minradius_preds, max_step_postds);
                }
 
                ~Shared ()
                {
                  mean_samples /= double(num_proc);
                  mean_truncations /= double(num_proc);
                  INFO ("mean number of samples per step = " + str (mean_samples));
                  if (mean_truncations) {
                    INFO ("mean number of steps between rejection sampling truncations = " + str (1.0/mean_truncations));
                    INFO ("maximum truncation error = " + str (max_max_truncation));
                  } else {
                    INFO ("no rejection sampling truncations occurred");
                  }
                }
 
                void update_stats (double mean_samples_per_run, double mean_truncations_per_run, double max_truncation) const
                {
                  mean_samples += mean_samples_per_run;
                  mean_truncations += mean_truncations_per_run;
                  if (max_truncation > max_max_truncation)
                    max_max_truncation = max_truncation;
                  ++num_proc;
                }
 
                float internal_step_size() const override { return step_size / float(num_samples); }
 
                size_t lmax, num_samples, max_trials;
                float sin_max_angle_ho, fod_power;
                Math::SH::PrecomputedAL<float> precomputer;
 
              private:
                mutable double mean_samples, mean_truncations, max_max_truncation;
                mutable int num_proc;
            };
 
 
 
 
 
 
 
 
 
            iFOD2 (const Shared& shared) :
              MethodBase (shared),
              S (shared),
              source (S.source),
              mean_sample_num (0),
              num_sample_runs (0),
              num_truncations (0),
              max_truncation (0.0),
              positions (S.num_samples),
              calib_positions (S.num_samples),
              tangents (S.num_samples),
              calib_tangents (S.num_samples),
              sample_idx (S.num_samples)
          {
            calibrate (*this);
          }
 
            iFOD2 (const iFOD2& that) :
              MethodBase (that.S),
              S (that.S),
              source (S.source),
              calibrate_ratio (that.calibrate_ratio),
              mean_sample_num (0),
              num_sample_runs (0),
              num_truncations (0),
              max_truncation (0.0),
              calibrate_list (that.calibrate_list),
              positions (S.num_samples),
              calib_positions (S.num_samples),
              tangents (S.num_samples),
              calib_tangents (S.num_samples),
              sample_idx (S.num_samples)
          {
          }
 
 
 
            ~iFOD2 ()
            {
              if (num_sample_runs)
                S.update_stats (calibrate_list.size() + float(mean_sample_num)/float(num_sample_runs),
                    float(num_truncations) / float(num_sample_runs),
                    max_truncation);
            }
 
 
 
 
            bool init() override
            {
              if (!get_data (source))
                return false;
 
              if (!S.init_dir.allFinite()) {
 
                const Eigen::Vector3f init_dir (dir);
 
                for (size_t n = 0; n < S.max_seed_attempts; n++) {
                  dir = init_dir.allFinite() ? rand_dir (init_dir) : random_direction();
                  half_log_prob0 = FOD (dir);
                  if (std::isfinite (half_log_prob0) && (half_log_prob0 > S.init_threshold))
                    goto end_init;
                }
 
              } else {
 
                dir = S.init_dir;
                half_log_prob0 = FOD (dir);
                if (std::isfinite (half_log_prob0) && (half_log_prob0 > S.init_threshold))
                  goto end_init;
 
              }
 
              return false;
 
end_init:
              half_log_prob0_seed = half_log_prob0 = 0.5 * std::log (half_log_prob0);
              sample_idx = S.num_samples; // Force arc calculation on first iteration
              return true;
            }
 
 
 
            term_t next () override
            {
 
              if (++sample_idx < S.num_samples) {
                pos = positions[sample_idx];
                dir = tangents [sample_idx];
                return CONTINUE;
              }
 
              Eigen::Vector3f next_pos, next_dir;
 
              float max_val = 0.0;
              for (size_t i = 0; i < calibrate_list.size(); ++i) {
                get_path (calib_positions, calib_tangents, rotate_direction (dir, calibrate_list[i]));
                float val = path_prob (calib_positions, calib_tangents);
                if (std::isnan (val))
                  return EXIT_IMAGE;
                else if (val > max_val)
                  max_val = val;
              }
 
              if (max_val <= 0.0)
                return CALIBRATOR;
 
              max_val *= calibrate_ratio;
 
              num_sample_runs++;
 
              for (size_t n = 0; n < S.max_trials; n++) {
                float val = rand_path_prob ();
 
                if (val > max_val) {
                  DEBUG ("max_val exceeded!!! (val = " + str(val) + ", max_val = " + str (max_val) + ")");
                  ++num_truncations;
                  if (val/max_val > max_truncation)
                    max_truncation = val/max_val;
                }
 
                if (uniform(rng) < val/max_val) {
                  mean_sample_num += n;
                  half_log_prob0 = last_half_log_probN;
                  pos = positions[0];
                  dir = tangents [0];
                  sample_idx = 0;
                  return CONTINUE;
                }
              }
 
              return MODEL;
            }
 
 
            float get_metric (const Eigen::Vector3f& position, const Eigen::Vector3f& direction) override
            {
              if (!get_data (source, position))
                return 0.0;
              return FOD (direction);
            }
 
 
            // Restore proper probability from the FOD at the track seed point
            void reverse_track() override
            {
              half_log_prob0 = half_log_prob0_seed;
              sample_idx = S.num_samples;
              MethodBase::reverse_track();
            }
 
 
            void truncate_track (GeneratedTrack& tck, const size_t length_to_revert_from, const size_t revert_step) override
            {
              // OK, if we know length_to_revert_from, we can reconstruct what sample_idx was at that point
              size_t sample_idx_at_full_length = (length_to_revert_from - tck.get_seed_index()) % S.num_samples;
              // Unfortunately can't distinguish between sample_idx = 0 and sample_idx = S.num_samples
              // However the former would result in zero truncation with revert_step = 1...
              if (!sample_idx_at_full_length)
                sample_idx_at_full_length = S.num_samples;
              const size_t points_to_remove = sample_idx_at_full_length + ((revert_step - 1) * S.num_samples);
              if (tck.get_seed_index() + points_to_remove >= tck.size()) {
                tck.clear();
                pos = { NaN, NaN, NaN };
                dir = { NaN, NaN, NaN };
                return;
              }
              const size_t new_size = length_to_revert_from - points_to_remove;
              if (tck.size() == 2 || new_size == 1)
                dir = (tck[1] - tck[0]).normalized();
              else if (new_size != tck.size())
                dir = (tck[new_size] - tck[new_size - 2]).normalized();
              tck.resize (new_size);
 
              // Need to get the path probability contribution from the FOD at this point
              pos = tck.back();
              get_data (source);
              half_log_prob0 = 0.5 * std::log (FOD (dir));
 
              // Make sure that arc is re-calculated when next() is called
              sample_idx = S.num_samples;
 
              // Need to update sgm_depth appropriately, remembering that it is tracked by exec
              if (S.is_act())
                act().sgm_depth = (act().sgm_depth > points_to_remove) ? act().sgm_depth - points_to_remove : 0;
            }
 
 
 
          private:
            const Shared& S;
            Interpolator<Image<float>>::type source;
            float calibrate_ratio, half_log_prob0, last_half_log_probN, half_log_prob0_seed;
            size_t mean_sample_num, num_sample_runs, num_truncations;
            float max_truncation;
            vector<Eigen::Vector3f> calibrate_list;
 
            // Store list of points in the currently-calculated arc
            vector<Eigen::Vector3f> positions, calib_positions;
            vector<Eigen::Vector3f> tangents, calib_tangents;
 
            // Generate an arc only when required, and on the majority of next() calls, simply return the next point
            //   in the arc - more dense structural image sampling
            size_t sample_idx;
 
 
 
            FORCE_INLINE float FOD (const Eigen::Vector3f& direction) const
            {
              return (S.precomputer ?
                  S.precomputer.value (values, direction) :
                  Math::SH::value (values, direction, S.lmax)
                  );
            }
 
            FORCE_INLINE float FOD (const Eigen::Vector3f& position, const Eigen::Vector3f& direction)
            {
              if (!get_data (source, position))
                return NaN;
              return FOD (direction);
            }
 
 
 
 
            FORCE_INLINE float rand_path_prob ()
            {
              get_path (positions, tangents, rand_dir (dir));
              return path_prob (positions, tangents);
            }
 
 
 
            float path_prob (vector<Eigen::Vector3f>& positions, vector<Eigen::Vector3f>& tangents)
            {
 
              // Early exit for ACT when path is not sensible
              if (S.is_act()) {
                if (!act().fetch_tissue_data (positions[S.num_samples - 1]))
                  return (NaN);
                if (act().tissues().get_csf() >= 0.5)
                  return 0.0;
              }
 
              float log_prob = half_log_prob0;
              for (size_t i = 0; i < S.num_samples; ++i) {
 
                float fod_amp = FOD (positions[i], tangents[i]);
                if (std::isnan (fod_amp))
                  return NaN;
                if (fod_amp < S.threshold)
                  return 0.0;
                fod_amp = std::log (fod_amp);
                if (i < S.num_samples-1) {
                  log_prob += fod_amp;
                } else {
                  last_half_log_probN = 0.5*fod_amp;
                  log_prob += last_half_log_probN;
                }
              }
 
              return std::exp (S.fod_power * log_prob);
            }
 
 
          protected:
            void get_path (vector<Eigen::Vector3f>& positions, vector<Eigen::Vector3f>& tangents, const Eigen::Vector3f& end_dir) const
            {
              float cos_theta = end_dir.dot (dir);
              cos_theta = std::min (cos_theta, float(1.0));
              float theta = std::acos (cos_theta);
 
              if (theta) {
 
                Eigen::Vector3f curv = end_dir - cos_theta * dir;
                curv.normalize();
                float R = S.step_size / theta;
 
                for (size_t i = 0; i < S.num_samples-1; ++i) {
                  float a = (theta * (i+1)) / S.num_samples;
                  float cos_a = std::cos (a);
                  float sin_a = std::sin (a);
                  positions[i] = pos + R * (sin_a * dir + (float(1.0) - cos_a) * curv);
                  tangents[i] = cos_a * dir + sin_a * curv;
                }
                positions[S.num_samples-1] = pos + R * (std::sin (theta) * dir + (float(1.0)-cos_theta) * curv);
                tangents[S.num_samples-1]  = end_dir;
 
              } else { // straight on:
 
                for (size_t i = 0; i < S.num_samples; ++i) {
                  float f = (i+1) * (S.step_size / S.num_samples);
                  positions[i] = pos + f * dir;
                  tangents[i]  = dir;
                }
 
              }
            }
 
 
 
            FORCE_INLINE Eigen::Vector3f rand_dir (const Eigen::Vector3f& d) { return (random_direction (d, S.max_angle_ho, S.sin_max_angle_ho)); }
 
 
 
          private:
            class Calibrate
            { MEMALIGN(Calibrate)
              public:
                Calibrate (iFOD2& method) :
                  P (method),
                  fod (P.values),
                  vox (P.S.vox()),
                  positions (P.S.num_samples),
                  tangents (P.S.num_samples) {
                    Math::SH::delta (fod, Eigen::Vector3f (0.0, 0.0, 1.0), P.S.lmax);
                    init_log_prob = 0.5 * std::log (Math::SH::value (P.values, Eigen::Vector3f (0.0, 0.0, 1.0), P.S.lmax));
                  }
 
                float operator() (float el)
                {
                  P.pos = { 0.0f, 0.0f, 0.0f };
                  P.get_path (positions, tangents, Eigen::Vector3f (std::sin (el), 0.0, std::cos(el)));
 
                  float log_prob = init_log_prob;
                  for (size_t i = 0; i < P.S.num_samples; ++i) {
                    float prob = Math::SH::value (P.values, tangents[i], P.S.lmax) * (1.0 - (positions[i][0] / vox));
                    if (prob <= 0.0)
                      return 0.0;
                    prob = std::log (prob);
                    if (i < P.S.num_samples-1)
                      log_prob += prob;
                    else
                      log_prob += 0.5*prob;
                  }
 
                  return std::exp (P.S.fod_power * log_prob);
                }
 
              private:
                iFOD2& P;
                Eigen::VectorXf& fod;
                const float vox;
                float init_log_prob;
                vector<Eigen::Vector3f> positions, tangents;
            };
 
            friend void calibrate<iFOD2> (iFOD2& method);
 
        };
 
 
 
      }
    }
  }
}
 
#endif