Merge pull request #2056 from SCIInstitute/disentangled_4d_ssm

Spatiotemporal SSM - Disentangled Approach

Author: akenmorris

.gitignore

@@ -35,6 +35,7 @@
 .DS_*
 .directory
 .idea/
+.vscode/
 
 # PYC files
 *.pyc

Libs/Optimize/CorrespondenceMode.h

@@ -8,6 +8,8 @@ namespace shapeworks {
     EnsembleRegressionEntropy = 3,
     EnsembleMixedEffectsEntropy = 4,
     MeshBasedGeneralEntropy = 5,
-    MeshBasedGeneralMeanEnergy = 6
+    MeshBasedGeneralMeanEnergy = 6,
+    DisentagledEnsembleEntropy = 7,
+    DisentangledEnsembleMeanEnergy = 8
   };
 }

Libs/Optimize/Function/DisentangledCorrespondenceFunction.cpp

@@ -0,0 +1,300 @@
+#include "DisentangledCorrespondenceFunction.h"
+#include <string>
+#include "Libs/Optimize/Domain/ImageDomainWithGradients.h"
+#include "Libs/Optimize/Utils/ParticleGaussianModeWriter.h"
+#include "Libs/Utils/Utils.h"
+#include <tbb/parallel_for.h>
+#include "vnl/algo/vnl_symmetric_eigensystem.h"
+
+namespace shapeworks {
+void DisentangledCorrespondenceFunction ::WriteModes(const std::string& prefix, int n) const {
+  typename ParticleGaussianModeWriter<VDimension>::Pointer writer = ParticleGaussianModeWriter<VDimension>::New();
+  writer->SetShapeMatrix(m_ShapeMatrix);
+  writer->SetFileName(prefix.c_str());
+  writer->SetNumberOfModes(n);
+  writer->Update();
+}
+
+void DisentangledCorrespondenceFunction::ComputeCovarianceMatrices() {
+  const unsigned int num_N = m_ShapeMatrix->cols(); // Total Number of subjects
+  const unsigned int num_T = m_ShapeMatrix->GetDomainsPerShape(); // Total Number of Time points
+  const unsigned int num_dims    = m_ShapeMatrix->rows() / num_T; // (dM X T) / T = dM
+  this->Initialize();
+  // computation across time cohort 
+  tbb::parallel_for(tbb::blocked_range<size_t>{0, num_T}, [&](const tbb::blocked_range<size_t>& r)
+  {
+    // Iterate t = 1....T
+    for (size_t time_inst = r.begin(); time_inst < r.end(); ++time_inst) {
+      // Build objective matrix Z
+      vnl_matrix_type z;
+      z.clear();
+      z.set_size(num_dims, num_N);
+      z.fill(0.0);
+      unsigned int row_idx_start = time_inst * num_dims;
+      z = m_ShapeMatrix->extract(num_dims, num_N, row_idx_start, 0);
+
+      // Resize Gradient Updates matrix for current time instance
+      if (m_Time_PointsUpdate->at(time_inst).rows() != num_dims || m_Time_PointsUpdate->at(time_inst).cols() != num_N)
+      {
+        m_Time_PointsUpdate->at(time_inst).set_size(num_dims, num_N);
+        m_Time_PointsUpdate->at(time_inst).fill(0.0);
+      }
+
+      // Compute mean and mean centred objective matrix for current time instance t_i
+      vnl_matrix_type points_minus_mean_t;
+      points_minus_mean_t.clear();
+      points_minus_mean_t.set_size(num_dims, num_N);
+      points_minus_mean_t.fill(0.0);
+      Eigen::MatrixXd inv_cov_t;
+      inv_cov_t.setZero();
+
+      m_points_mean_time_cohort->at(time_inst).clear();
+      m_points_mean_time_cohort->at(time_inst).set_size(num_dims, 1);
+
+      for (unsigned int j = 0; j < num_dims; ++j)
+      {
+        double sum_across_col = 0.0;
+        for (unsigned int i = 0; i < num_N; ++i)
+        {
+          sum_across_col += z(j, i);
+        }
+        m_points_mean_time_cohort->at(time_inst).put(j,0, sum_across_col/(double)num_N);
+      }
+
+      for (unsigned int j = 0; j < num_dims; j++)
+      {
+        for (unsigned int i = 0; i < num_N; i++)
+        {
+          points_minus_mean_t(j, i) = z(j, i) - m_points_mean_time_cohort->at(time_inst).get(j,0);
+        }
+      }
+
+      vnl_diag_matrix<double> W_t;
+      vnl_matrix_type gramMat_t(num_N, num_N, 0.0); // gram matrix = Y^T X Y
+      vnl_matrix_type pinvMat_t(num_N, num_N, 0.0); // inverse of gram Matrix
+
+      if (this->m_UseMeanEnergy)
+      {
+        pinvMat_t.set_identity();
+        m_InverseCovMatrices_time_cohort->at(time_inst).setZero();
+      }
+      else
+      {
+        gramMat_t = points_minus_mean_t.transpose()* points_minus_mean_t;
+        vnl_svd <double> svd(gramMat_t);
+        vnl_matrix_type UG = svd.U();
+        W_t = svd.W();
+        vnl_diag_matrix<double> invLambda_t = svd.W();
+        invLambda_t.set_diagonal(invLambda_t.get_diagonal()/(double)(num_N-1) + m_MinimumVariance);
+        invLambda_t.invert_in_place();
+
+        pinvMat_t = (UG * invLambda_t) * UG.transpose();
+        vnl_matrix_type projMat_t = points_minus_mean_t * UG;
+        const auto lhs = projMat_t * invLambda_t;
+        const auto rhs = invLambda_t * projMat_t.transpose();
+        inv_cov_t.resize(num_dims, num_dims);
+        Utils::multiply_into(inv_cov_t, lhs, rhs);
+      }
+      // Update Gradient points update infor
+      m_Time_PointsUpdate->at(time_inst).update(points_minus_mean_t * pinvMat_t);
+      double currentEnergy_t = 0.0;
+      if (m_UseMeanEnergy) currentEnergy_t = points_minus_mean_t.frobenius_norm();
+      else
+      {
+        m_MinimumEigenValue_time_cohort[time_inst] = W_t(0)*W_t(0) + m_MinimumVariance;
+        for (unsigned int i = 0; i < num_N; i++)
+        {
+          double val_i = W_t(i)*W_t(i) + m_MinimumVariance;
+          if ( val_i < m_MinimumEigenValue_time_cohort[time_inst])
+            m_MinimumEigenValue_time_cohort[time_inst] = val_i;
+          currentEnergy_t += log(val_i);
+        }
+      }
+      currentEnergy_t /= 2.0;
+      if (m_UseMeanEnergy) m_MinimumEigenValue_time_cohort[time_inst] = currentEnergy_t / 2.0;
+      // Update Inv Covariance Matrix
+      m_InverseCovMatrices_time_cohort->at(time_inst) = inv_cov_t;
+    }
+  });
+
+  // computation across shape cohort 
+  tbb::parallel_for(tbb::blocked_range<size_t>{0, num_N}, [&](const tbb::blocked_range<size_t>& r)
+  {
+    // Iterate n = 1....N
+    for (size_t sub = r.begin(); sub < r.end(); ++sub) {
+      // Build objective matrix Z
+      vnl_matrix_type z;
+      z.clear();
+      z.set_size(num_dims, num_N);
+      z.fill(0.0);
+      for(unsigned int t = 0; t < num_T; ++t){
+        unsigned int row_start = num_dims * t;
+        vnl_matrix_type time_vec = m_ShapeMatrix->extract(num_dims, 1, row_start, sub);
+        z.set_columns(t, time_vec);
+      }
+
+      // Resize Gradient Updates matrix for current time instance
+      if (m_Shape_PointsUpdate->at(sub).rows() != num_dims || m_Shape_PointsUpdate->at(sub).cols() != num_T)
+      {
+        m_Shape_PointsUpdate->at(sub).set_size(num_dims, num_T);
+        m_Shape_PointsUpdate->at(sub).fill(0.0);
+      }
+
+      // Compute mean and mean centred objective matrix for current time instance t_i
+      vnl_matrix_type points_minus_mean_n;
+      points_minus_mean_n.clear();
+      points_minus_mean_n.set_size(num_dims, num_T);
+      points_minus_mean_n.fill(0.0);
+      Eigen::MatrixXd inv_cov_n;
+      inv_cov_n.setZero();
+
+      m_points_mean_shape_cohort->at(sub).clear();
+      m_points_mean_shape_cohort->at(sub).set_size(num_dims, 1);
+
+      for (unsigned int j = 0; j < num_dims; ++j)
+      {
+        double sum_across_col = 0.0;
+        for (unsigned int i = 0; i < num_T; ++i)
+        {
+          sum_across_col += z(j, i);
+        }
+        m_points_mean_shape_cohort->at(sub).put(j,0, sum_across_col/(double)num_T);
+      }
+
+      for (unsigned int j = 0; j < num_dims; j++)
+      {
+        for (unsigned int i = 0; i < num_T; i++)
+        {
+          points_minus_mean_n(j, i) = z(j, i) - m_points_mean_shape_cohort->at(sub).get(j,0);
+        }
+      }
+
+      vnl_diag_matrix<double> W_n;
+      vnl_matrix_type gramMat_n(num_T, num_T, 0.0); // gram matrix = Y^T X Y
+      vnl_matrix_type pinvMat_n(num_T, num_T, 0.0); // inverse of gram Matrix
+
+      if (this->m_UseMeanEnergy)
+      {
+        pinvMat_n.set_identity();
+        m_InverseCovMatrices_shape_cohort->at(sub).setZero();
+      }
+      else
+      {
+        gramMat_n = points_minus_mean_n.transpose() * points_minus_mean_n;
+        vnl_svd <double> svd(gramMat_n);
+        vnl_matrix_type UG = svd.U();
+        W_n = svd.W();
+        vnl_diag_matrix<double> invLambda_n = svd.W();
+        invLambda_n.set_diagonal(invLambda_n.get_diagonal()/(double)(num_T-1) + m_MinimumVariance);
+        invLambda_n.invert_in_place();
+
+        pinvMat_n = (UG * invLambda_n) * UG.transpose();
+        vnl_matrix_type projMat_n = points_minus_mean_n * UG;
+        const auto lhs = projMat_n * invLambda_n;
+        const auto rhs = invLambda_n * projMat_n.transpose();
+        inv_cov_n.resize(num_dims, num_dims);
+        Utils::multiply_into(inv_cov_n, lhs, rhs);
+      }
+
+      // Update Gradient points update infor
+      m_Shape_PointsUpdate->at(sub).update(points_minus_mean_n * pinvMat_n);
+      double currentEnergy_n = 0.0;
+      if (m_UseMeanEnergy) currentEnergy_n = points_minus_mean_n.frobenius_norm();
+      else
+      {
+        m_MinimumEigenValue_shape_cohort[sub] = W_n(0)*W_n(0) + m_MinimumVariance;
+        for (unsigned int i = 0; i < num_T; i++)
+        {
+          double val_i = W_n(i) * W_n(i) + m_MinimumVariance;
+          if (val_i < m_MinimumEigenValue_shape_cohort[sub])
+            m_MinimumEigenValue_shape_cohort[sub] = val_i;
+          currentEnergy_n += log(val_i);
+        }
+      }
+      currentEnergy_n /= 2.0;
+      if (m_UseMeanEnergy) m_MinimumEigenValue_shape_cohort[sub] = currentEnergy_n / 2.0;
+      // Update Inv Covariance Matrix
+      m_InverseCovMatrices_shape_cohort->at(sub) = inv_cov_n;
+    }
+  });
+
+}
+
+DisentangledCorrespondenceFunction::VectorType DisentangledCorrespondenceFunction ::Evaluate(unsigned int idx, unsigned int d,
+                                                                                       const ParticleSystem* system,
+                                                                                       double& maxdt,
+                                                                                       double& energy) const {
+
+  const unsigned int num_N = m_ShapeMatrix->cols(); // Total number of subjects
+  const unsigned int num_T = m_ShapeMatrix->GetDomainsPerShape(); // Total number of time points
+    
+  const unsigned int cur_sub = d / num_T; // index of current subject
+  const unsigned int cur_time_point = d % num_T;
+
+  // maximum update possible = sum of max possible updates across both cohorts (time and shape)
+  maxdt  = m_MinimumEigenValue_shape_cohort[cur_sub] + m_MinimumEigenValue_time_cohort[cur_time_point];
+
+  VectorType gradE; // gradient update vector for Point defined by ParticleSystem system of domain d and dimension index idx
+  unsigned int shape_matrix_start_idx = 0;
+
+  int dom = d % num_T;
+  for (int i = 0; i < dom; i++)
+    shape_matrix_start_idx += system->GetNumberOfParticles(i) * VDimension;
+  shape_matrix_start_idx += idx*VDimension;
+
+  unsigned int particle_idx = VDimension * idx;
+
+  // Energy computation across time cohort
+  vnl_matrix_type Xi_time_cohort(3,1,0.0);
+  Xi_time_cohort(0,0) = m_ShapeMatrix->operator()(shape_matrix_start_idx  , cur_sub) - m_points_mean_time_cohort->at(cur_time_point).get(particle_idx, 0);
+  Xi_time_cohort(1,0) = m_ShapeMatrix->operator()(shape_matrix_start_idx+1, cur_sub) - m_points_mean_time_cohort->at(cur_time_point).get(particle_idx+1, 0);
+  Xi_time_cohort(2,0) = m_ShapeMatrix->operator()(shape_matrix_start_idx+2, cur_sub) - m_points_mean_time_cohort->at(cur_time_point).get(particle_idx+2, 0);
+  vnl_matrix_type tmp1_time(3, 3, 0.0);
+  if (this->m_UseMeanEnergy) {
+    tmp1_time.set_identity();
+  } else {
+    Eigen::MatrixXd region = m_InverseCovMatrices_time_cohort->at(cur_time_point).block(particle_idx, particle_idx, 3, 3);
+    // convert to vnl
+    for (unsigned int i = 0; i < 3; i++) {
+      for (unsigned int j = 0; j < 3; j++) {
+        tmp1_time(i, j) = region(i, j);
+      }
+    }
+  }
+  vnl_matrix_type tmp_time = Xi_time_cohort.transpose()*tmp1_time;
+  tmp_time *= Xi_time_cohort;
+
+  // Energy computation across shape cohort
+  vnl_matrix_type Xi_shape_cohort(3,1,0.0);
+  Xi_shape_cohort(0,0) = m_ShapeMatrix->operator()(shape_matrix_start_idx  , cur_sub) - m_points_mean_shape_cohort->at(cur_sub).get(particle_idx, 0);
+  Xi_shape_cohort(1,0) = m_ShapeMatrix->operator()(shape_matrix_start_idx+1, cur_sub) - m_points_mean_shape_cohort->at(cur_sub).get(particle_idx+1, 0);
+  Xi_shape_cohort(2,0) = m_ShapeMatrix->operator()(shape_matrix_start_idx+2, cur_sub) - m_points_mean_shape_cohort->at(cur_sub).get(particle_idx+2, 0);
+  vnl_matrix_type tmp1_shape(3, 3, 0.0);
+  if (this->m_UseMeanEnergy) {
+    tmp1_shape.set_identity();
+  } else {
+    Eigen::MatrixXd region = m_InverseCovMatrices_shape_cohort->at(cur_sub).block(particle_idx, particle_idx, 3, 3);
+    // convert to vnl
+    for (unsigned int i = 0; i < 3; i++) {
+      for (unsigned int j = 0; j < 3; j++) {
+        tmp1_time(i, j) = region(i, j);
+      }
+    }
+  }
+  vnl_matrix_type tmp_shape = Xi_shape_cohort.transpose()*tmp1_shape;
+  tmp_shape *= Xi_shape_cohort;
+
+  // Net Energy
+  energy = tmp_time(0,0) + tmp_shape(0, 0); 
+
+  // Net Gradient
+  for (unsigned int i = 0; i< VDimension; i++)
+  {
+    gradE[i] = m_Time_PointsUpdate->at(cur_time_point).get(particle_idx + i, cur_sub) + m_Shape_PointsUpdate->at(cur_sub).get(particle_idx + i, cur_time_point);
+  }
+  return system->TransformVector(gradE,
+                                   system->GetInversePrefixTransform(d) *
+                                   system->GetInverseTransform(d));
+}
+
+}  // namespace shapeworks