CPU-based gradient descent (#111)

* Add gradient descent type option

* Add CPU based gradient descent

* Add gradient descent option to settings

* Add gradient descent settings to readme

* Update README.md

README.md

@@ -36,7 +36,10 @@ cmake --build build --config Release --target install
   - Defaults to random. Optional: Use another data set as the initial embedding coordinates, e.g. the first two [PCA](https://github.com/ManiVaultStudio/PcaPlugin/) components.
   - Defaults to rescaling the initial coordinates such that the first embedding dimension has a standard deviation of 0.0001. If turned off, the random initialization will uniformly sample coordinates from a circle with radius 1.
   - See e.g. [The art of using t-SNE for single-cell transcriptomics](https://doi.org/10.1038/s41467-019-13056-x) for more details on recommended t-SNE settings
-- Changed t-SNE gradient descent parameters are not taken into account when "continuing" the gradient descent, but when "reinitializing" they are
-- knn settings specify search structure construction and query characteristics:
+- Gradient Descent:
+  - GPU-based implementation (default) requires OpenGL 3.3 and benefits from compute shaders (introduced in OpenGL 4.4 and not available on Apple devices)
+  - CPU-based implementation of [Barnes-Hut t-SNE](https://jmlr.org/papers/v15/vandermaaten14a.html) automatically sets θ to `min(0.5, max(0.0, (numPoints - 1000.0) * 0.00005))`
+  - Changes to gradient descent parameters are not taken into account when "continuing" the gradient descent, but when "reinitializing" they are
+- kNN (specify search structure construction and query characteristics):
   - (Annoy) Trees & Checks: correspond to `n_trees` and `search_k`, see their [docs](https://github.com/spotify/annoy?tab=readme-ov-file#tradeoffs)
   - (HNSW): M & ef: are detailed in the respective [docs](https://github.com/nmslib/hnswlib/blob/master/ALGO_PARAMS.md#hnsw-algorithm-parameters)

src/Common/GradientDescentSettingsAction.cpp

@@ -9,11 +9,13 @@ GradientDescentSettingsAction::GradientDescentSettingsAction(QObject* parent, Ts
     _tsneParameters(tsneParameters),
     _exaggerationFactorAction(this, "Exaggeration factor"),
     _exaggerationIterAction(this, "Exaggeration iterations"),
-    _exponentialDecayAction(this, "Exponential decay")
+    _exponentialDecayAction(this, "Exponential decay"),
+    _gradienDescentTypeAction(this, "GD implementation")
 {
     addAction(&_exaggerationFactorAction);
     addAction(&_exaggerationIterAction);
     addAction(&_exponentialDecayAction);
+    addAction(&_gradienDescentTypeAction);
 
     _exaggerationFactorAction.setDefaultWidgetFlags(IntegralAction::SpinBox);
     _exaggerationIterAction.setDefaultWidgetFlags(IntegralAction::SpinBox);
@@ -23,8 +25,11 @@ GradientDescentSettingsAction::GradientDescentSettingsAction(QObject* parent, Ts
     _exaggerationIterAction.initialize(1, 10000, 250);
     _exponentialDecayAction.initialize(1, 10000, 70);
 
+    _gradienDescentTypeAction.initialize({ "GPU", "CPU" });
+
     _exaggerationFactorAction.setToolTip("Defaults to 4 + number of points / 60'000");
     _exponentialDecayAction.setToolTip("Iterations after 'Exaggeration iterations' during \nwhich the exaggeration factor exponentionally decays towards 1");
+    _gradienDescentTypeAction.setToolTip("Gradient Descent Implementation: GPU (A-tSNE), CPU (Barnes-Hut)");
 
     const auto updateExaggerationFactor = [this]() -> void {
         _tsneParameters.setExaggerationFactor(_exaggerationFactorAction.getValue());
@@ -38,12 +43,22 @@ GradientDescentSettingsAction::GradientDescentSettingsAction(QObject* parent, Ts
         _tsneParameters.setExponentialDecayIter(_exponentialDecayAction.getValue());
     };
 
+    const auto updateGradienDescentTypeAction = [this]() -> void {
+        switch (_gradienDescentTypeAction.getCurrentIndex())
+        {
+        case 0: _tsneParameters.setGradienDescentType(GradienDescentType::GPU); break;
+        case 1: _tsneParameters.setGradienDescentType(GradienDescentType::CPU); break;
+        }
+
+    };
+
     const auto updateReadOnly = [this]() -> void {
         const auto enable = !isReadOnly();
 
         _exaggerationFactorAction.setEnabled(enable);
         _exaggerationIterAction.setEnabled(enable);
         _exponentialDecayAction.setEnabled(enable);
+        _gradienDescentTypeAction.setEnabled(enable);
     };
 
     connect(&_exaggerationFactorAction, &DecimalAction::valueChanged, this, [this, updateExaggerationFactor](const float value) {
@@ -58,13 +73,18 @@ GradientDescentSettingsAction::GradientDescentSettingsAction(QObject* parent, Ts
         updateExponentialDecay();
     });
 
+    connect(&_gradienDescentTypeAction, &OptionAction::currentIndexChanged, this, [this, updateGradienDescentTypeAction](const std::int32_t& currentIndex) {
+        updateGradienDescentTypeAction();
+    });
+
     connect(this, &GroupAction::readOnlyChanged, this, [this, updateReadOnly](const bool& readOnly) {
         updateReadOnly();
     });
 
     updateExaggerationFactor();
     updateExaggerationIter();
     updateExponentialDecay();
+    updateGradienDescentTypeAction();
     updateReadOnly();
 }
 
@@ -75,6 +95,7 @@ void GradientDescentSettingsAction::fromVariantMap(const QVariantMap& variantMap
     _exaggerationFactorAction.fromParentVariantMap(variantMap);
     _exaggerationIterAction.fromParentVariantMap(variantMap);
     _exponentialDecayAction.fromParentVariantMap(variantMap);
+    _gradienDescentTypeAction.fromParentVariantMap(variantMap);
 }
 
 QVariantMap GradientDescentSettingsAction::toVariantMap() const
@@ -84,6 +105,7 @@ QVariantMap GradientDescentSettingsAction::toVariantMap() const
     _exaggerationFactorAction.insertIntoVariantMap(variantMap);
     _exaggerationIterAction.insertIntoVariantMap(variantMap);
     _exponentialDecayAction.insertIntoVariantMap(variantMap);
+    _gradienDescentTypeAction.insertIntoVariantMap(variantMap);
 
     return variantMap;
 }

src/Common/GradientDescentSettingsAction.h

@@ -3,6 +3,7 @@
 #include "actions/DecimalAction.h"
 #include "actions/GroupAction.h"
 #include "actions/IntegralAction.h"
+#include "actions/OptionAction.h"
 
 using namespace mv::gui;
 
@@ -30,6 +31,7 @@ class GradientDescentSettingsAction : public GroupAction
     DecimalAction& getExaggerationFactorAction() { return _exaggerationFactorAction; };
     IntegralAction& getExaggerationIterAction() { return _exaggerationIterAction; };
     IntegralAction& getExponentialDecayAction() { return _exponentialDecayAction; };
+    OptionAction& getGadienDescentTypeAction() { return _gradienDescentTypeAction; };
 
 public: // Serialization
 
@@ -50,4 +52,5 @@ class GradientDescentSettingsAction : public GroupAction
     DecimalAction           _exaggerationFactorAction;  /** Exaggeration factor action */
     IntegralAction          _exaggerationIterAction;    /** Exaggeration iteration action */
     IntegralAction          _exponentialDecayAction;    /** Exponential decay action */
+    OptionAction            _gradienDescentTypeAction;  /** GPU or CPU gradient descent */
 };

src/Common/TsneAnalysis.cpp

@@ -23,6 +23,7 @@ TsneWorker::TsneWorker(TsneParameters tsneParameters) :
     _probabilityDistribution(),
     _hasProbabilityDistribution(false),
     _GPGPU_tSNE(),
+    _CPU_tSNE(),
     _embedding(),
     _outEmbedding(),
     _offscreenBuffer(nullptr),
@@ -203,7 +204,7 @@ void TsneWorker::computeSimilarities()
     }
 
     qDebug() << "================================================================================";
-    qDebug() << "A-tSNE: Computed probability distribution: " << t / 1000 << " seconds";
+    qDebug() << "tSNE: Computed probability distribution: " << t / 1000 << " seconds";
     qDebug() << "--------------------------------------------------------------------------------";
 
     _tasks->getComputingSimilaritiesTask().setFinished();
@@ -219,36 +220,79 @@ void TsneWorker::computeGradientDescent(uint32_t iterations)
         emit embeddingUpdate(tsneData);
         };
 
-    _tasks->getInitializeTsneTask().setRunning();
+    auto initGPUTSNE = [this]() {
+        // Initialize offscreen buffer
+        double t_buffer = 0.0;
+        {
+            hdi::utils::ScopedTimer<double> timer(t_buffer);
+            _offscreenBuffer->bindContext();
+        }
+        qDebug() << "tSNE: Set up offscreen buffer in " << t_buffer / 1000 << " seconds.";
 
-    // Initialize offscreen buffer
-    double t_buffer = 0.0;
-    {
-        hdi::utils::ScopedTimer<double> timer(t_buffer);
-        _offscreenBuffer->bindContext();
-    }
-    qDebug() << "A-tSNE: Set up offscreen buffer in " << t_buffer / 1000 << " seconds.";
+        if (!_GPGPU_tSNE.isInitialized())
+        {
+            auto params = tsneParameters();
 
-    // Initialize GPGPU-SNE
-    double t_init = 0.0;
-    if (!_GPGPU_tSNE.isInitialized())
-    {
-        hdi::utils::ScopedTimer<double> timer(t_init);
+            // In case of HSNE, the _probabilityDistribution is a non-summetric transition matrix and initialize() symmetrizes it here
+            if (_hasProbabilityDistribution)
+                _GPGPU_tSNE.initialize(_probabilityDistribution, &_embedding, params);
+            else
+                _GPGPU_tSNE.initializeWithJointProbabilityDistribution(_probabilityDistribution, &_embedding, params);
+
+            qDebug() << "A-tSNE (GPU): Exaggeration factor: " << params._exaggeration_factor << ", exaggeration iterations: " << params._remove_exaggeration_iter << ", exaggeration decay iter: " << params._exponential_decay_iter;
+        }
+    };
+
+    auto initCPUTSNE = [this]() {
+        if (!_CPU_tSNE.isInitialized())
+        {
+            auto params = tsneParameters();
+
+            double theta = std::min(0.5, std::max(0.0, (_numPoints - 1000.0) * 0.00005));
+            _CPU_tSNE.setTheta(theta);
+
+            // In case of HSNE, the _probabilityDistribution is a non-summetric transition matrix and initialize() symmetrizes it here
+            if (_hasProbabilityDistribution)
+                _CPU_tSNE.initialize(_probabilityDistribution, &_embedding, params);
+            else
+                _CPU_tSNE.initializeWithJointProbabilityDistribution(_probabilityDistribution, &_embedding, params);
 
-        auto params = tsneParameters();
+            qDebug() << "t-SNE (CPU, Barnes-Hut): Exaggeration factor: " << params._exaggeration_factor << ", exaggeration iterations: " << params._remove_exaggeration_iter << ", exaggeration decay iter: " << params._exponential_decay_iter << ", theta: " << theta;
+        }
+    };
+
+    auto initTSNE = [this, initGPUTSNE, initCPUTSNE, updateEmbedding]() {
+        double t_init = 0.0;
+        {
+            hdi::utils::ScopedTimer<double> timer(t_init);
+
+            if (_tsneParameters.getGradienDescentType() == GradienDescentType::GPU)
+                initGPUTSNE();
+            else
+                initCPUTSNE();
+
+            updateEmbedding(_outEmbedding);
+        }
+        qDebug() << "tSNE: Init t-SNE " << t_init / 1000 << " seconds.";
+    };
 
-        // In case of HSNE, the _probabilityDistribution is a non-summetric transition matrix and initialize() symmetrizes it here
-        if (_hasProbabilityDistribution)
-            _GPGPU_tSNE.initialize(_probabilityDistribution, &_embedding, params);
+    auto singleTSNEIteration = [this]() {
+        if (_tsneParameters.getGradienDescentType() == GradienDescentType::GPU)
+            _GPGPU_tSNE.doAnIteration();
         else
-            _GPGPU_tSNE.initializeWithJointProbabilityDistribution(_probabilityDistribution, &_embedding, params);
+            _CPU_tSNE.doAnIteration();
+    };
 
-        qDebug() << "A-tSNE: Exaggeration factor: " << params._exaggeration_factor << ", exaggeration iterations: " << params._remove_exaggeration_iter << ", exaggeration decay iter: " << params._exponential_decay_iter;
-    }
+    auto gradientDescentCleanup = [this]() {
+        if (_tsneParameters.getGradienDescentType() == GradienDescentType::GPU)
+            _offscreenBuffer->releaseContext();
+        else
+            return; // Nothing to do for CPU implementation
+    };
 
-    updateEmbedding(_outEmbedding);
+    _tasks->getInitializeTsneTask().setRunning();
 
-    qDebug() << "A-tSNE: Init t-SNE " << t_init / 1000 << " seconds.";
+    initTSNE();
 
     _tasks->getInitializeTsneTask().setFinished();
 
@@ -258,7 +302,7 @@ void TsneWorker::computeGradientDescent(uint32_t iterations)
     double elapsed = 0;
     double t_grad = 0;
     {
-        qDebug() << "A-tSNE: Computing " << endIteration - beginIteration << " gradient descent iterations...";
+        qDebug() << "tSNE: Computing " << endIteration - beginIteration << " gradient descent iterations...";
 
         _tasks->getComputeGradientDescentTask().setRunning();
         _tasks->getComputeGradientDescentTask().setSubtasks(iterations);
@@ -272,8 +316,8 @@ void TsneWorker::computeGradientDescent(uint32_t iterations)
 
             hdi::utils::ScopedTimer<double> timer(t_grad);
 
-            // Perform a GPGPU-SNE iteration
-            _GPGPU_tSNE.doAnIteration();
+            // Perform t-SNE iteration
+            singleTSNEIteration();
 
             if (_currentIteration > 0 && _tsneParameters.getUpdateCore() > 0 && _currentIteration % _tsneParameters.getUpdateCore() == 0)
                 updateEmbedding(_outEmbedding);
@@ -294,15 +338,15 @@ void TsneWorker::computeGradientDescent(uint32_t iterations)
             QCoreApplication::processEvents();
         }
 
-        _offscreenBuffer->releaseContext();
+        gradientDescentCleanup();
 
         updateEmbedding(_outEmbedding);
 
         _tasks->getComputeGradientDescentTask().setFinished();
     }
 
     qDebug() << "--------------------------------------------------------------------------------";
-    qDebug() << "A-tSNE: Finished embedding in: " << elapsed / 1000 << " seconds, with " << _currentIteration << " total iterations (" << endIteration - beginIteration << " new iterations)";
+    qDebug() << "tSNE: Finished embedding in: " << elapsed / 1000 << " seconds, with " << _currentIteration << " total iterations (" << endIteration - beginIteration << " new iterations)";
     qDebug() << "================================================================================";
 
     emit finished();

src/Common/TsneAnalysis.h

@@ -6,6 +6,7 @@
 
 #include "hdi/dimensionality_reduction/gradient_descent_tsne_texture.h"
 #include "hdi/dimensionality_reduction/hd_joint_probability_generator.h"
+#include "hdi/dimensionality_reduction/sparse_tsne_user_def_probabilities.h"
 #include "hdi/dimensionality_reduction/tsne_parameters.h"
 
 #include <Task.h>
@@ -41,6 +42,9 @@ class TsneWorker : public QObject
 {
     Q_OBJECT
 
+    using GradienDescentGPU = hdi::dr::GradientDescentTSNETexture;
+    using GradienDescentCPU = hdi::dr::SparseTSNEUserDefProbabilities<float>;
+
 private:
     // default construction is inaccessible to outsiders
     TsneWorker(TsneParameters tsneParameters);
@@ -97,7 +101,8 @@ public slots:
     std::vector<float>                      _data;                          /** High-dimensional input data */
     ProbDistMatrix                          _probabilityDistribution;       /** High-dimensional probability distribution encoding point similarities */
     bool                                    _hasProbabilityDistribution;    /** Check if the worker was initialized with a probability distribution or data */
-    hdi::dr::GradientDescentTSNETexture     _GPGPU_tSNE;                    /** GPGPU t-SNE gradient descent implementation */
+    GradienDescentGPU                       _GPGPU_tSNE;                    /** GPGPU t-SNE gradient descent implementation */
+    GradienDescentCPU                       _CPU_tSNE;                      /** CPU t-SNE gradient descent implementation */
     hdi::data::Embedding<float>             _embedding;                     /** Storage of current embedding */
     TsneData                                _outEmbedding;                  /** Transfer embedding data array */
     OffscreenBuffer*                        _offscreenBuffer;               /** Offscreen OpenGL buffer required to run the gradient descent */

src/Common/TsneParameters.h

@@ -1,5 +1,12 @@
 #pragma once
 
+enum class GradienDescentType
+{
+    GPU,
+    CPU,
+};
+
+
 class TsneParameters
 {
 public:
@@ -11,7 +18,8 @@ class TsneParameters
         _numDimensionsOutput(2),
         _presetEmbedding(false),
         _exaggerationFactor(4),
-        _updateCore(10)
+        _updateCore(10),
+        _gradienDescentType(GradienDescentType::GPU)
     {
 
     }
@@ -23,6 +31,7 @@ class TsneParameters
     void setNumDimensionsOutput(int numDimensionsOutput) { _numDimensionsOutput = numDimensionsOutput; }
     void setPresetEmbedding(bool presetEmbedding) { _presetEmbedding = presetEmbedding; }
     void setExaggerationFactor(double exaggerationFactor) { _exaggerationFactor = exaggerationFactor; }
+    void setGradienDescentType(GradienDescentType gradienDescentType) { _gradienDescentType = gradienDescentType; }
     void setUpdateCore(int updateCore) { _updateCore = updateCore; }
 
     int getNumIterations() const { return _numIterations; }
@@ -32,6 +41,7 @@ class TsneParameters
     int getNumDimensionsOutput() const { return _numDimensionsOutput; }
     int getPresetEmbedding() const { return _presetEmbedding; }
     int getExaggerationFactor() const { return _exaggerationFactor; }
+    GradienDescentType getGradienDescentType() const { return _gradienDescentType; }
     int getUpdateCore() const { return _updateCore; }
 
 private:
@@ -42,6 +52,7 @@ class TsneParameters
     int _numDimensionsOutput;
     double _exaggerationFactor;
     bool _presetEmbedding;
+    GradienDescentType _gradienDescentType;     // Whether to use CPU or GPU gradient descent
 
     int _updateCore;        // Gradient descent iterations after which the embedding data set in ManiVault's core will be updated
 };