fixed calculation of energy

cmd/dirflip.cpp

@@ -82,7 +82,7 @@ class Shared { MEMALIGN(Shared)
       vector3_type b = { directions(j,0), directions(j,1), directions(j,2) };
       if (signs[i] < 0) a = -a;
       if (signs[j] < 0) b = -b;
-      return 1.0 / (a-b).squaredNorm();
+      return 1.0 / (a-b).norm();
     }
 
 

cmd/dirgen.cpp

@@ -15,12 +15,14 @@
 #include "command.h"
 #include "progressbar.h"
 #include "math/rng.h"
+#include "thread.h"
 #include "math/gradient_descent.h"
 #include "math/check_gradient.h"
 #include "dwi/directions/file.h"
 
-#define DEFAULT_POWER 2
+#define DEFAULT_POWER 1
 #define DEFAULT_NITER 10000
+#define DEFAULT_RESTARTS 10
 
 
 using namespace MR;
@@ -29,33 +31,44 @@ using namespace App;
 void usage ()
 {
 
-AUTHOR = "J-Donald Tournier (jdtournier@gmail.com)";
+  AUTHOR = "J-Donald Tournier (jdtournier@gmail.com)";
 
-SYNOPSIS = "Generate a set of uniformly distributed directions using a bipolar electrostatic repulsion model";
+  SYNOPSIS = "Generate a set of uniformly distributed directions using a bipolar electrostatic repulsion model";
 
-REFERENCES 
-  + "Jones, D.; Horsfield, M. & Simmons, A. "
-  "Optimal strategies for measuring diffusion in anisotropic systems by magnetic resonance imaging. "
-  "Magnetic Resonance in Medicine, 1999, 42: 515-525"
-
-  + "Papadakis, N. G.; Murrills, C. D.; Hall, L. D.; Huang, C. L.-H. & Adrian Carpenter, T. "
-  "Minimal gradient encoding for robust estimation of diffusion anisotropy. "
-  "Magnetic Resonance Imaging, 2000, 18: 671-679";
+  DESCRIPTION 
+    + "Directions are distributed by analogy to an electrostatic repulsion system, with each direction "
+    "corresponding to a single electrostatic charge (for -unipolar), or a pair of diametrically opposed charges "
+    "for the default bipolar case). The energy of the system is determined based on the Coulomb repulstion, "
+    "which assumed the form 1/r^power, where r is the distance between any pair of charges, and p is the power "
+    "assumed for the repulsion law (default: 1). The minimum energy state is obtained by gradient descent.";
 
-ARGUMENTS
-  + Argument ("ndir", "the number of directions to generate.").type_integer (6, std::numeric_limits<int>::max())
-  + Argument ("dirs", "the text file to write the directions to, as [ az el ] pairs.").type_file_out();
 
-OPTIONS
-  + Option ("power", "specify exponent to use for repulsion power law (default: " + str(DEFAULT_POWER) + "). This must be a power of 2 (i.e. 2, 4, 8, 16, ...).")
-  +   Argument ("exp").type_integer (2, std::numeric_limits<int>::max())
+  REFERENCES 
+    + "Jones, D.; Horsfield, M. & Simmons, A. "
+    "Optimal strategies for measuring diffusion in anisotropic systems by magnetic resonance imaging. "
+    "Magnetic Resonance in Medicine, 1999, 42: 515-525"
 
-  + Option ("niter", "specify the maximum number of iterations to perform (default: " + str(DEFAULT_NITER) + ").")
-  +   Argument ("num").type_integer (1, std::numeric_limits<int>::max())
+    + "Papadakis, N. G.; Murrills, C. D.; Hall, L. D.; Huang, C. L.-H. & Adrian Carpenter, T. "
+    "Minimal gradient encoding for robust estimation of diffusion anisotropy. "
+    "Magnetic Resonance Imaging, 2000, 18: 671-679";
 
-  + Option ("unipolar", "optimise assuming a unipolar electrostatic repulsion model rather than the bipolar model normally assumed in DWI")
+  ARGUMENTS
+    + Argument ("ndir", "the number of directions to generate.").type_integer (6, std::numeric_limits<int>::max())
+    + Argument ("dirs", "the text file to write the directions to, as [ az el ] pairs.").type_file_out();
 
-  + Option ("cartesian", "Output the directions in Cartesian coordinates [x y z] instead of [az el].");
+  OPTIONS
+    + Option ("power", "specify exponent to use for repulsion power law (default: " + str(DEFAULT_POWER) + "). This must be a power of 2 (i.e. 1, 2, 4, 8, 16, ...).")
+    +   Argument ("exp").type_integer (1, std::numeric_limits<int>::max())
+
+    + Option ("niter", "specify the maximum number of iterations to perform (default: " + str(DEFAULT_NITER) + ").")
+    +   Argument ("num").type_integer (1, std::numeric_limits<int>::max())
+
+    + Option ("restarts", "specify the number of restarts to perform (default: " + str(DEFAULT_RESTARTS) + ").")
+    +   Argument ("num").type_integer (1, std::numeric_limits<int>::max())
+
+    + Option ("unipolar", "optimise assuming a unipolar electrostatic repulsion model rather than the bipolar model normally assumed in DWI")
+
+    + Option ("cartesian", "Output the directions in Cartesian coordinates [x y z] instead of [az el].");
 
 }
 
@@ -79,52 +92,70 @@ class ProjectedUpdate { MEMALIGN(ProjectedUpdate)
 
 
 
+
+
+
+
 class Energy { MEMALIGN(Energy)
   public:
-    Energy (size_t n_directions, int power, bool bipolar, const Eigen::VectorXd& init_directions) : 
-      ndir (n_directions),
-      power (power),
-      bipolar (bipolar),
-      init_directions (init_directions) { }
+    Energy (ProgressBar& progress) : 
+      progress (progress),
+      ndirs (to<int> (argument[0])),
+      bipolar (!(get_options ("unipolar").size())),
+      power (0),
+      directions (3 * ndirs) { }
 
     FORCE_INLINE double fast_pow (double x, int p) {
       return p == 1 ? x : fast_pow (x*x, p/2);
     }
 
     using value_type = double;
 
-    size_t size () const { return 3 * ndir; }
+    size_t size () const { return 3 * ndirs; }
 
     // set x to original directions provided in constructor. 
     // The idea is to save the directions from one run to initialise next run
     // at higher power.
-    double init (Eigen::VectorXd& x) {
-      x = init_directions;
+    double init (Eigen::VectorXd& x) 
+    {
+      Math::RNG::Normal<double> rng;
+      for (ssize_t n = 0; n < ndirs; ++n) {
+        auto d = x.segment (3*n,3);
+        d[0] = rng();
+        d[1] = rng();
+        d[2] = rng();
+        d.normalize();
+      }
       return 0.01;
     }
 
+
+
+    // function executed by optimiser at each iteration:
     double operator() (const Eigen::VectorXd& x, Eigen::VectorXd& g) {
       double E = 0.0;
       g.setZero();
 
-      for (size_t i = 0; i < ndir-1; ++i) {
+      for (size_t i = 0; i < ndirs-1; ++i) {
         auto d1 = x.segment (3*i, 3);
         auto g1 = g.segment (3*i, 3);
-        for (size_t j = i+1; j < ndir; ++j) {
+        for (size_t j = i+1; j < ndirs; ++j) {
           auto d2 = x.segment (3*j, 3);
           auto g2 = g.segment (3*j, 3);
 
           Eigen::Vector3d r = d1-d2;
           double _1_r2 = 1.0 / r.squaredNorm();
-          double e = fast_pow (_1_r2, power/2); 
+          double _1_r = std::sqrt (_1_r2);
+          double e = fast_pow (_1_r, power); 
           E += e;
           g1 -= (power * e * _1_r2) * r; 
           g2 += (power * e * _1_r2) * r; 
 
           if (bipolar) {
             r = d1+d2;
             _1_r2 = 1.0 / r.squaredNorm();
-            e = fast_pow (_1_r2, power/2); 
+            _1_r = std::sqrt (_1_r2);
+            e = fast_pow (_1_r, power); 
             E += e;
             g1 -= (power * e * _1_r2) * r; 
             g2 -= (power * e * _1_r2) * r; 
@@ -134,77 +165,102 @@ class Energy { MEMALIGN(Energy)
       }
 
       // constrain gradients to lie tangent to unit sphere:
-      for (size_t n = 0; n < ndir; ++n) 
+      for (size_t n = 0; n < ndirs; ++n) 
         g.segment(3*n,3) -= x.segment(3*n,3).dot (g.segment(3*n,3)) * x.segment(3*n,3);
 
       return E;
     }
 
 
-  protected:
-    size_t ndir;
-    int power;
-    bool bipolar;
-    const Eigen::VectorXd& init_directions;
-};
 
+    // function executed per thread:
+    void execute () 
+    {
+      size_t this_start = 0;
+      while ((this_start = current_start.fetch_add (1)) < restarts) {
+        INFO ("launching start " + str (this_start));
+        double E = 0.0;
 
+        for (power = 1; power <= target_power; power *= 2) {
+          Math::GradientDescent<Energy,ProjectedUpdate> optim (*this, ProjectedUpdate());
 
-void run () {
-  size_t niter = get_option_value ("niter", DEFAULT_NITER);
-  int target_power = get_option_value ("power", DEFAULT_POWER);
-  bool bipolar = !(get_options ("unipolar").size());
-  int ndirs = to<int> (argument[0]);
+          INFO ("start " + str(this_start) + ": setting power = " + str (power));
+          optim.init();
 
-  Eigen::VectorXd directions (3*ndirs);
+          size_t iter = 0;
+          for (; iter < niter; iter++) {
+            if (!optim.iterate()) 
+              break;
 
-  { // random initialisation:
-    Math::RNG::Normal<double> rng;
-    for (ssize_t n = 0; n < ndirs; ++n) {
-      auto d = directions.segment (3*n,3);
-      d[0] = rng();
-      d[1] = rng();
-      d[2] = rng();
-      d.normalize();
+            DEBUG ("start " + str(this_start) + ": [ " + str (iter) + " ] (pow = " + str (power) + ") E = " + str (optim.value(), 8)
+                + ", grad = " + str (optim.gradient_norm(), 8));
+
+            std::lock_guard<std::mutex> lock (mutex);
+            ++progress;
+          }
+
+          directions = optim.state();
+          E = optim.value();
+        }
+
+
+
+        std::lock_guard<std::mutex> lock (mutex);
+        if (E < best_E) {
+          best_E = E;
+          best_directions = directions;
+        }
+      }
     }
-  }
 
-  // optimisation proper:
-  {
-    ProgressBar progress ("Optimising directions");
-    for (int power = 2; power <= target_power; power *= 2) {
-      Energy energy (ndirs, power, bipolar, directions);
 
-      Math::GradientDescent<Energy,ProjectedUpdate> optim (energy, ProjectedUpdate());
+    static size_t restarts;
+    static int target_power;
+    static size_t niter;
+    static double best_E;
+    static Eigen::VectorXd best_directions;
 
-      INFO ("setting power = " + str (power));
-      optim.init();
+  protected:
+    ProgressBar& progress;
+    size_t ndirs;
+    bool bipolar;
+    int power;
+    Eigen::VectorXd directions;
+    double E;
+
+    static std::mutex mutex;
+    static std::atomic<size_t> current_start;
+};
 
-      //Math::check_function_gradient (energy, optim.state(), 0.001);
-      //return;
 
-      size_t iter = 0;
-      for (; iter < niter; iter++) {
-        if (!optim.iterate()) 
-          break;
+size_t Energy::restarts (DEFAULT_RESTARTS);
+int Energy::target_power (DEFAULT_POWER);
+size_t Energy::niter (DEFAULT_NITER);
+std::mutex Energy::mutex;
+std::atomic<size_t> Energy::current_start (0);
+double Energy::best_E = std::numeric_limits<double>::infinity();
+Eigen::VectorXd Energy::best_directions;
 
-        INFO ("[ " + str (iter) + " ] (pow = " + str (power) + ") E = " + str (optim.value(), 8)
-            + ", grad = " + str (optim.gradient_norm(), 8));
 
-        progress.update ([&]() { return "Optimising directions (power " + str(power) 
-            + ", energy: " + str(optim.value(), 8) + ", gradient: " + str(optim.gradient_norm(), 8) + ", iteration " + str(iter) + ")"; });
-      }
 
-      directions = optim.state();
 
-      progress.set_text ("Optimising directions (power " + str(power) 
-        + ", energy: " + str(optim.value(), 8) + ", gradient: " + str(optim.gradient_norm(), 8) + ", iteration " + str(iter) + ")");
-    }
+void run () 
+{
+  Energy::restarts = get_option_value ("restarts", DEFAULT_RESTARTS);
+  Energy::target_power = get_option_value ("power", DEFAULT_POWER);
+  Energy::niter = get_option_value ("niter", DEFAULT_NITER);
+
+  { 
+    ProgressBar progress ("Optimising directions up to power " + str(Energy::target_power) + " (" + str(Energy::restarts) + " restarts)");
+    Energy energy_functor (progress);
+    auto threads = Thread::run (Thread::multi (energy_functor), "energy function");
   }
 
+  CONSOLE ("final energy = " + str(Energy::best_E));
+  size_t ndirs = Energy::best_directions.size()/3;
   Eigen::MatrixXd directions_matrix (ndirs, 3);
   for (int n = 0; n < ndirs; ++n) 
-    directions_matrix.row (n) = directions.segment (3*n, 3);
+    directions_matrix.row (n) = Energy::best_directions.segment (3*n, 3);
 
   DWI::Directions::save (directions_matrix, argument[1], get_options ("cartesian").size());
 }

cmd/dirmerge.cpp

@@ -38,12 +38,17 @@ ARGUMENTS
   + Argument ("out", "the output directions file, with each row listing "
       "the X Y Z gradient directions, the b-value, and an index representing "
       "the phase encode direction").type_file_out();
+
+OPTIONS
+  + Option ("unipolar_weight", 
+      "set the weight given to the unipolar electrostatic repulsion model compared to the "
+      "bipolar model (default: 0.2).");
 }
 
 
 
 using value_type = double;
-using Direction = std::array<value_type,3>;
+using Direction = Eigen::Matrix<value_type,3,1>;
 using DirectionSet = vector<Direction>;
 
 
@@ -62,6 +67,8 @@ inline std::ostream& operator<< (std::ostream& stream, const OutDir& d) {
 void run () 
 {
   size_t num_subsets = argument[0];
+  value_type unipolar_weight = App::get_option_value ("unipolar_weight", 0.2);
+  value_type bipolar_weight = 1.0 - unipolar_weight;
 
 
   vector<vector<DirectionSet>> dirs;
@@ -80,7 +87,7 @@ void run ()
       auto m = DWI::Directions::load_cartesian (argument[current++]);
       DirectionSet set;
       for (ssize_t r = 0; r < m.rows(); ++r)
-        set.push_back ({ { m(r,0), m(r,1), m(r,2) } });
+        set.push_back (Direction (m(r,0), m(r,1), m(r,2)));
       d.push_back (set);
     }
     INFO ("found b = " + str(bvalue[nb]) + ", " + 
@@ -104,24 +111,18 @@ void run ()
 
   auto push = [&](size_t b, size_t p, size_t n) 
   { 
-    merged.push_back ({ { { dirs[b][p][n][0], dirs[b][p][n][1], dirs[b][p][n][2] } }, b, p }); 
+    merged.push_back ({ Direction (dirs[b][p][n][0], dirs[b][p][n][1], dirs[b][p][n][2]), b, p }); 
     dirs[b][p].erase (dirs[b][p].begin()+n); 
   };
 
-  auto energy_pair = [](const Direction& a, const Direction& b) 
+  auto energy_pair = [&](const Direction& a, const Direction& b) 
   {
     // use combination of mono- and bi-polar electrostatic repulsion models 
     // to ensure adequate coverage of eddy-current space as well as 
     // orientation space. Use a moderate bias, favouring the bipolar model.
-    return 1.2 / (
-        Math::pow2 (b[0] - a[0]) + 
-        Math::pow2 (b[1] - a[1]) + 
-        Math::pow2 (b[2] - a[2]) 
-        ) + 1.0 / (
-        Math::pow2 (b[0] + a[0]) + 
-        Math::pow2 (b[1] + a[1]) + 
-        Math::pow2 (b[2] + a[2]) 
-        );
+
+    return (unipolar_weight+bipolar_weight) / (b-a).norm()
+         + bipolar_weight / (b+a).norm();
   };
 
   auto energy = [&](size_t b, size_t p, size_t n)