Large cohort scaling improvements

Libs/Optimize/Function/CorrespondenceFunction.cpp

@@ -8,6 +8,8 @@
 #include "vnl/algo/vnl_svd.h"
 #include "vnl/vnl_diag_matrix.h"
 
+#include <Eigen/Dense>
+
 namespace shapeworks {
 
 void CorrespondenceFunction::ComputeUpdates(const ParticleSystem* c) {
@@ -178,48 +180,54 @@ void CorrespondenceFunction::ComputeUpdates(const ParticleSystem* c) {
 
   vnl_diag_matrix<double> W;
 
-  vnl_matrix_type gramMat(num_samples, num_samples, 0.0);
-  vnl_matrix_type pinvMat(num_samples, num_samples, 0.0);  // gramMat inverse
+  vnl_matrix_type Q;
 
   if (this->m_UseMeanEnergy) {
-    pinvMat.set_identity();
+    // In mean-energy mode the inverse covariance is the identity, so Q is just
+    // the mean-centered points. Skip building an N×N identity and the multiply,
+    // which is a no-op costing O(P·N²) per iteration.
+    Q = points_minus_mean;
   } else {
+    using RowMatrixXd = Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>;
+
     TIME_START("correspondence::gramMat");
 
     // Create Eigen maps that point to the VNL data
-    Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>> points_minus_mean_map(
-        points_minus_mean.data_block(), points_minus_mean.rows(), points_minus_mean.cols());
+    Eigen::Map<RowMatrixXd> points_minus_mean_map(points_minus_mean.data_block(), points_minus_mean.rows(),
+                                                  points_minus_mean.cols());
 
-    // Perform the computation directly on the mapped data
+    // gramMat = Yᵀ Y is symmetric positive-semidefinite.
     Eigen::MatrixXd gramMat_eigen = points_minus_mean_map.transpose() * points_minus_mean_map;
 
-    // Copy result back to VNL matrix
-    gramMat.set_size(gramMat_eigen.rows(), gramMat_eigen.cols());
-    std::memcpy(gramMat.data_block(), gramMat_eigen.data(), gramMat_eigen.size() * sizeof(double));
-
     TIME_STOP("correspondence::gramMat");
 
     TIME_START("correspondence::svd");
-    vnl_svd<double> svd(gramMat);
-    vnl_matrix_type UG = svd.U();
-    W = svd.W();
-    vnl_diag_matrix<double> invLambda = svd.W();
+    // The Gram matrix is symmetric positive-semidefinite, so a self-adjoint
+    // eigensolver yields the same eigenpairs as a general SVD but much faster,
+    // and it parallelizes the surrounding dense algebra through Eigen.
+    Eigen::SelfAdjointEigenSolver<Eigen::MatrixXd> eig(gramMat_eigen);
+    const Eigen::VectorXd& eigvals = eig.eigenvalues();
+    const Eigen::MatrixXd& UG = eig.eigenvectors();
     TIME_STOP("correspondence::svd");
 
-    invLambda.set_diagonal(invLambda.get_diagonal() / (double)(num_samples - 1) + m_MinimumVariance);
-    invLambda.invert_in_place();
+    // Energy term uses the eigenvalues (== singular values of the Gram matrix).
+    W = vnl_diag_matrix<double>(num_samples);
+    for (int i = 0; i < num_samples; ++i) W(i) = eigvals(i);
+
+    // pinvMat = UG · diag(1 / (λ/(N-1) + α)) · UGᵀ  (symmetric).
+    Eigen::ArrayXd invLambda = (eigvals.array() / (double)(num_samples - 1) + m_MinimumVariance).inverse();
 
     TIME_START("correspondence::pinvMat");
-    pinvMat = (UG * invLambda) * UG.transpose();
+    Eigen::MatrixXd pinv = UG * invLambda.matrix().asDiagonal() * UG.transpose();
     TIME_STOP("correspondence::pinvMat");
 
-    // Note: The full dM × dM inverse covariance matrix is NOT needed.
-    // Per thesis equation 2.35, the gradient is simply: Y × (Y^T Y + αI)^{-1}
-    // which is Y × pinvMat. No covariance_multiply required.
+    // Q = Y · pinvMat  (thesis eq. 2.35: gradient is Y × (YᵀY + αI)⁻¹).
+    RowMatrixXd Q_eigen;
+    Q_eigen.noalias() = points_minus_mean_map * pinv;
+    Q.set_size(num_dims, num_samples);
+    std::memcpy(Q.data_block(), Q_eigen.data(), Q_eigen.size() * sizeof(double));
   }
 
-  vnl_matrix_type Q = points_minus_mean * pinvMat;
-
   // Compute the update matrix in coordinate space by multiplication with the
   // Jacobian.  Each shape gradient must be transformed by a different Jacobian
   // so we have to do this individually for each shape (sample).

Libs/Optimize/Function/LegacyCorrespondenceFunction.cpp

@@ -1,6 +1,7 @@
 
 #include "LegacyCorrespondenceFunction.h"
 
+#include <cstring>
 #include <string>
 
 #include "Libs/Optimize/Domain/ImageDomainWithGradients.h"
@@ -160,31 +161,41 @@ void LegacyCorrespondenceFunction ::ComputeCovarianceMatrix() {
 
   vnl_diag_matrix<double> W;
 
-  vnl_matrix_type gramMat(num_samples, num_samples, 0.0);
-  vnl_matrix_type pinvMat(num_samples, num_samples, 0.0);  // gramMat inverse
-
   if (this->m_UseMeanEnergy) {
-    pinvMat.set_identity();
+    // In mean-energy mode the inverse covariance is the identity, so the update
+    // is just the mean-centered points. Skip building an N×N identity and the
+    // subsequent multiply, which is a no-op costing O(P·N²) per iteration.
+    m_PointsUpdate->update(points_minus_mean);
   } else {
-    gramMat = points_minus_mean.transpose() * points_minus_mean;
-
-    vnl_svd<double> svd(gramMat);
-
-    vnl_matrix_type UG = svd.U();
-    W = svd.W();
-
-    vnl_diag_matrix<double> invLambda = svd.W();
-
-    invLambda.set_diagonal(invLambda.get_diagonal() / (double)(num_samples - 1) + m_MinimumVariance);
-    invLambda.invert_in_place();
-
-    pinvMat = (UG * invLambda) * UG.transpose();
-
-    // Note: The full dM × dM inverse covariance matrix is NOT needed.
-    // Per thesis equation 2.35, the gradient is simply: Y × (Y^T Y + αI)^{-1}
-    // which is Y × pinvMat. No covariance matrix multiplication required.
+    vnl_matrix_type gramMat = points_minus_mean.transpose() * points_minus_mean;
+
+    // gramMat = Yᵀ Y is symmetric positive-semidefinite, so a self-adjoint
+    // eigensolver yields the same eigenpairs as a general SVD but much faster,
+    // and it parallelizes the surrounding dense algebra through Eigen. gramMat is
+    // symmetric, so the row-major VNL buffer maps to an identical Eigen matrix.
+    Eigen::Map<Eigen::MatrixXd> gram_map(gramMat.data_block(), num_samples, num_samples);
+    Eigen::SelfAdjointEigenSolver<Eigen::MatrixXd> eig(gram_map);
+    const Eigen::VectorXd& eigvals = eig.eigenvalues();
+    const Eigen::MatrixXd& UG = eig.eigenvectors();
+
+    // Energy term uses the eigenvalues (== singular values of the Gram matrix).
+    W = vnl_diag_matrix<double>(num_samples);
+    for (unsigned int i = 0; i < num_samples; ++i) W(i) = eigvals(i);
+
+    // pinvMat = UG · diag(1 / (λ/(N-1) + α)) · UGᵀ  (symmetric).
+    using RowMatrixXd = Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>;
+    Eigen::ArrayXd invLambda = (eigvals.array() / (double)(num_samples - 1) + m_MinimumVariance).inverse();
+    Eigen::MatrixXd pinv = UG * invLambda.matrix().asDiagonal() * UG.transpose();
+
+    // update = Y · pinvMat  (thesis eq. 2.35: gradient is Y × (YᵀY + αI)⁻¹).
+    Eigen::Map<RowMatrixXd> Y_map(points_minus_mean.data_block(), num_dims, num_samples);
+    RowMatrixXd update_eigen;
+    update_eigen.noalias() = Y_map * pinv;
+
+    vnl_matrix_type update(num_dims, num_samples);
+    std::memcpy(update.data_block(), update_eigen.data(), update_eigen.size() * sizeof(double));
+    m_PointsUpdate->update(update);
   }
-  m_PointsUpdate->update(points_minus_mean * pinvMat);
 
   //     std::cout << m_PointsUpdate.extract(num_dims, num_samples,0,0) << std::endl;
 

docs/workflow/optimize-scaling.md

@@ -1,117 +1,117 @@
 # Optimization Performance & Scaling
 
-This page explains how `shapeworks optimize` scales with the size of your dataset: how runtime grows
-with the number of shapes and particles, and why very large cohorts tend to use only a single CPU
-core. If you are optimizing hundreds or thousands of shapes and seeing long runtimes or low CPU
-utilization, this page describes what to expect and what you can do about it.
+This page explains how `shapeworks optimize` scales with the size of your dataset, and what recent
+improvements changed. If you are optimizing hundreds or thousands of shapes and want to know what to
+expect for runtime and CPU usage, start here.
 
 ## What to expect
 
-- Runtime grows cubically with the number of shapes (subjects, N). Doubling the number of shapes can
-  multiply optimization time by roughly 8x in the large-cohort regime.
-- Large cohorts use about one CPU core. Small and medium datasets parallelize across many cores. As
-  the number of shapes grows into the thousands, the optimizer spends most of its time in a serial
-  step, so utilization falls toward a single core.
-- Particle count is much cheaper than shape count. Increasing the number of particles grows runtime
-  far more slowly than increasing the number of shapes.
-
-## What you can do
-
-- Use incremental optimization for large cohorts. It optimizes a small initial subset of shapes from
-  scratch, then adds the remaining shapes in batches. The expensive particle initialization runs only
-  on the initial subset, and each later batch starts from already-placed particles, so the
-  full-cohort passes need far fewer iterations. This lowers total runtime. It does not remove the
-  per-iteration cost: the final batch still includes every shape and pays the O(N³) SVD each
-  iteration. See [Incremental Supershapes](../use-cases/multistep/incremental_supershapes.md) for a
-  worked example. Incremental optimization is available through the Python and command-line
-  interfaces; ShapeWorks Studio does not currently support it.
-- Lowering the iteration count does not reduce the per-iteration cost. The expensive step runs every
-  iteration, so lowering `optimization_iterations` reduces the number of iterations but not the cost
-  of each one.
-- If you need higher-resolution models, adding particles is much cheaper than adding shapes.
-
-## Why
+- Runtime grows faster with the number of shapes (subjects, N) than with the number of particles. The
+  correspondence step at the center of each iteration decomposes an N×N matrix built from the shapes,
+  which costs O(N³), so adding shapes is much more expensive than adding particles.
+- Recent improvements cut this cost and restored multi-core use in both phases of a run. A full run at
+  N=1024 is about 3.5x faster than before, and the gap grows with N. See
+  [Recent improvements](#recent-improvements) below.
+- Adding particles is much cheaper than adding shapes.
 
-Each optimization iteration has two parts:
-
-| part | complexity | parallel? |
-|---|---|---|
-| per-particle gradient update (`tbb::parallel_for` over subjects) | O(N·P) | yes, across subjects |
-| correspondence entropy: an N×N Gram matrix and one SVD | O(N³) | no, single-threaded |
+## Recent improvements
 
-N is the number of subjects and P the number of particles. The per-particle gradient is spread across
-cores, but the correspondence step runs a serial SVD on an N×N matrix once per iteration, which is
-O(N³). As N grows, that serial cubic step dominates, so total runtime grows cubically and CPU usage
-falls toward one core. This happens whether or not `use_normals` is enabled, because both
-correspondence formulations build the same N×N SVD.
+Two changes to the per-iteration correspondence update made both phases of a run faster and brought
+back multi-core usage on large cohorts. Neither changes the optimization result.
 
-## Measurements
+1. Initialization no longer multiplies by an identity matrix. In the initialization (mean-energy)
+   phase the update reduces to the mean-centered points, but the code used to build an N×N identity
+   and run the full multiply every iteration, an O(P·N²) no-op. Removing it makes initialization far
+   faster at scale (the whole initialization phase at N=2048 went from about 10 minutes to about 45
+   seconds) and keeps it multi-core.
 
-Measured on a 24-core machine with 128 GB of memory. Cohorts were built by replicating a single femur
-mesh (15,002 vertices) N times with a small random rigid transform and per-vertex noise, so N can be
-varied over a wide range. Shape content does not affect the O(N³) SVD, which depends only on the
-matrix size. Metrics are reported per iteration, since a full 100-iteration run at N=4096 takes days.
+2. The optimization phase uses a symmetric eigensolver. The N×N matrix is symmetric
+   positive-semidefinite, so a symmetric eigensolver replaces the general SVD and the surrounding
+   matrix work runs multithreaded. This is faster per iteration by an amount that grows with N: about
+   2.6x at N=512, 8x at N=2048, and 14x at N=4096.
 
-### Runtime per iteration vs number of shapes
+End to end, both changes together (initialization plus optimization), on the same cohorts:
 
-![runtime vs N](../img/optimize-scaling/runtime_vs_N.png)
+![end to end before and after](../img/optimize-scaling/improvement_e2e_walltime.png)
 
-The curve is flat at small N, where the parallel work dominates, then bends up to follow the N³
-reference beyond N=512. The lines converge because the SVD cost does not depend on the particle
-count.
+| N | before | after | speedup |
+|---|---|---|---|
+| 256 | 10.0 s | 8.0 s | 1.3x |
+| 512 | 53.6 s | 22.6 s | 2.4x |
+| 1024 | 224.9 s | 64.0 s | 3.5x |
 
-### The serial SVD overtakes the parallel work
+Below a few hundred shapes the difference is small, because the parts these changes target only
+dominate at larger N. The payoff grows with the cohort.
 
-![phase vs N](../img/optimize-scaling/phase_vs_N.png)
+CPU utilization also recovers. Mean cores busy for each phase (of 24), before and after, at P=256:
 
-The serial SVD (slope near 3) passes the parallel per-particle work (slope near 1) at about N=1024 and
-dominates beyond that.
+| N | initialization before / after | optimization before / after |
+|---|---|---|
+| 256 | 11.9 / 13.8 | 9.0 / 10.9 |
+| 1024 | 4.0 / 12.7 | 4.0 / 7.4 |
+| 2048 | 2.0 / 12.6 | 2.4 / 7.0 |
 
-### CPU utilization vs number of shapes
+![CPU utilization vs N](../img/optimize-scaling/cores_vs_N.png)
 
-![cores vs N](../img/optimize-scaling/cores_vs_N.png)
+Before the changes, optimization utilization fell to about 1.5 cores at N=4096. After, it stays around
+6, because the eigensolve itself is still serial even though the matrix work around it is now
+threaded. So utilization is much better but still declines at the very high end.
 
-Mean cores busy rises to a peak around N=64, then falls as the serial SVD takes over, dropping to
-about 1.3 cores at N=4096. A cohort of thousands of shapes sits on the right side of this curve, in
-the single-core regime. The high utilization on small datasets is the left side of the same curve.
+## Why shapes cost more than particles
 
-### Particle scaling is much milder
+Each optimization iteration has two parts:
 
-![runtime vs P](../img/optimize-scaling/runtime_vs_P.png)
+| part | complexity | parallel? |
+|---|---|---|
+| per-particle gradient update (`tbb::parallel_for` over subjects) | O(N·P) | yes, across subjects |
+| correspondence: build the N×N matrix and decompose it | O(N³) | the eigensolve is serial; the surrounding matrix work runs multithreaded |
 
-Adding particles grows the cheaper terms, not the cubic SVD, so it scales far more slowly than adding
-shapes.
+N is the number of subjects and P the number of particles. The decomposition is O(N³) in the number
+of shapes, so runtime still grows cubically with N even after the improvements. The changes lower the
+constant and parallelize most of the surrounding work, but the eigensolve itself is serial, which is
+why very large cohorts remain expensive. This holds whether or not `use_normals` is enabled, because
+both correspondence formulations build the same N×N decomposition.
 
-## Initialization vs optimization
+![runtime vs N](../img/optimize-scaling/runtime_vs_N.png)
 
-A run has two phases. Initialization adds particles by splitting, running `iterations_per_split`
-iterations at each level. Optimization then runs `optimization_iterations` at the full particle count.
-The expensive SVD runs in the optimization phase; initialization uses a cheaper mean-energy mode. For
-a typical run, initialization is often the majority of the wall time.
+Optimization time per iteration, before and after. Both follow the N³ reference at large N, so the
+slope is unchanged, but the eigensolver change shifts the curve down by a widening margin, reaching
+about 14x at N=4096. At very small N the two are within measurement noise; the eigensolver has a small
+fixed overhead that does not matter once the matrices are tiny.
 
-![cores init vs opt](../img/optimize-scaling/cores_init_vs_opt.png)
+![runtime vs P](../img/optimize-scaling/runtime_vs_P.png)
 
-Core utilization for each phase measured in isolation (P=256):
+Adding particles grows the cheaper terms, not the cubic decomposition, so it scales far more slowly
+than adding shapes.
 
-| N | initialization cores | optimization cores |
-|---|---|---|
-| 256 | 11.9 | 9.0 |
-| 1024 | 4.0 | 4.0 |
-| 2048 | 2.0 | 2.4 |
+## What you can do for very large cohorts
 
-Initialization parallelizes well for typical cohorts of up to a few hundred shapes, but at thousands
-of shapes it also becomes serial-bound. Very large cohorts see reduced CPU utilization in both phases.
+- Use incremental optimization. It optimizes a small initial subset of shapes from scratch, then adds
+  the remaining shapes in batches. The expensive particle initialization runs only on the initial
+  subset, and each later batch starts from already-placed particles, so the full-cohort passes need
+  far fewer iterations. It does not remove the per-iteration cost: the final batch still includes
+  every shape and pays the O(N³) decomposition each iteration. See
+  [Incremental Supershapes](../use-cases/multistep/incremental_supershapes.md) for a worked example.
+  Incremental optimization is available through the Python and command-line interfaces; ShapeWorks
+  Studio does not currently support it.
+- Lowering the iteration count does not reduce the per-iteration cost. The decomposition runs every
+  iteration, so lowering `optimization_iterations` reduces the number of iterations but not the cost
+  of each one.
+- If you need higher-resolution models, adding particles is much cheaper than adding shapes.
 
-![cores mixed](../img/optimize-scaling/cores_mixed_vs_N.png)
+## Measurements
 
-The blended view of a realistic run shows high utilization at small N, falling toward a single core as
-the number of shapes reaches the thousands.
+Measured on a 24-core machine with 128 GB of memory. Cohorts were built by replicating a single femur
+mesh (15,002 vertices) N times with a small random rigid transform and per-vertex noise, so N can be
+varied over a wide range. Shape content does not affect the decomposition cost, which depends only on
+the matrix size. Before and after numbers come from back-to-back runs with the binary swapped per N
+under the same load, since machine contention alone can move cross-run timings by roughly 1.8x.
 
 ## Summary
 
-- The cubic runtime and single-core behavior on large cohorts are expected. A serial O(N³) SVD runs
-  once per optimization iteration.
+- Runtime still grows cubically with the number of shapes, because the correspondence step decomposes
+  an N×N matrix each iteration. Recent improvements lowered the constant and restored multi-core use,
+  but did not change the cubic shape.
 - The number of shapes drives the cost. The number of particles has less of an impact.
-- No setting removes the per-iteration O(N³) cost. Incremental optimization is the most useful option
-  for large cohorts: it does the expensive initialization on a small subset and lets later batches run
-  with far fewer iterations.
+- For very large cohorts, incremental optimization is the most useful option: it does the expensive
+  initialization on a small subset and lets later batches run with far fewer iterations.