Skip to content

Commit 8062b4f

Browse files
akenmorrisakenmorris
authored
Merge pull request #2483 from SCIInstitute/fixed_shape_space
Optimize correspondence function, add fixed domain precomputation
2 parents 82dbe62 + 8cdceef commit 8062b4f

6 files changed

Lines changed: 661 additions & 94 deletions

File tree

‎Libs/Optimize/Function/CorrespondenceFunction.cpp‎

Lines changed: 290 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11

22
#include "CorrespondenceFunction.h"
33

4-
#include <Libs/Utils/PlatformUtils.h>
54
#include <math.h>
65

76
#include "Libs/Utils/Utils.h"
@@ -12,7 +11,6 @@
1211
namespace shapeworks {
1312

1413
void CorrespondenceFunction::ComputeUpdates(const ParticleSystem* c) {
15-
TIME_SCOPE("CorrespondenceFunction::ComputeUpdates");
1614
num_dims = m_ShapeData->rows();
1715
num_samples = m_ShapeData->cols();
1816

@@ -26,6 +24,147 @@ void CorrespondenceFunction::ComputeUpdates(const ParticleSystem* c) {
2624

2725
m_PointsUpdate->fill(0.0);
2826

27+
// Lazy precomputation: do it on first call when shape data is populated
28+
if (m_NeedsPrecomputation && !m_HasPrecomputedFixedDomains && num_samples > 0) {
29+
m_NeedsPrecomputation = false;
30+
PrecomputeForFixedDomains(c);
31+
}
32+
33+
// =========================================================================
34+
// FIXED SHAPE SPACE PATH: Use precomputed data if available
35+
// =========================================================================
36+
if (m_HasPrecomputedFixedDomains && !m_UseMeanEnergy) {
37+
TIME_SCOPE("ComputeUpdates::FixedShapeSpace");
38+
39+
// Use precomputed mean from fixed shapes
40+
m_points_mean = m_FixedMean;
41+
42+
// For each non-fixed shape, compute its gradient using precomputed data
43+
vnl_matrix_type Q(num_dims, num_samples);
44+
Q.fill(0.0);
45+
46+
// Compute centered data for all shapes using fixed mean
47+
vnl_matrix_type points_minus_mean(num_dims, num_samples);
48+
for (int j = 0; j < num_samples; j++) {
49+
for (int i = 0; i < num_dims; i++) {
50+
points_minus_mean(i, j) = m_ShapeData->get(i, j) - (*m_FixedMean)(i, 0);
51+
}
52+
}
53+
54+
// Create Eigen maps for precomputed matrices (VNL uses row-major storage)
55+
Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>> Y_fixed_map(
56+
m_FixedY->data_block(), num_dims, m_PrecomputedNumFixedSamples);
57+
Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>> GradientBasis_map(
58+
m_FixedGradientBasis->data_block(), num_dims, m_PrecomputedNumFixedSamples);
59+
Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>> pinvMat_map(
60+
m_FixedPinvMat->data_block(), m_PrecomputedNumFixedSamples, m_PrecomputedNumFixedSamples);
61+
62+
// Regularization denominator
63+
const double N = m_PrecomputedNumFixedSamples - 1;
64+
const double N_sq = N * N;
65+
66+
// For each non-fixed shape, compute gradient using Schur complement
67+
for (int j = 0; j < num_samples; j++) {
68+
// Check if this shape is fixed (skip gradient computation for fixed shapes)
69+
bool is_fixed = false;
70+
for (int idx : m_FixedDomainIndices) {
71+
if (idx == j) {
72+
is_fixed = true;
73+
break;
74+
}
75+
}
76+
77+
if (!is_fixed) {
78+
// Extract y_j from points_minus_mean column (VNL is row-major, so column is strided)
79+
Eigen::VectorXd y_j(num_dims);
80+
for (int i = 0; i < num_dims; i++) {
81+
y_j[i] = points_minus_mean(i, j);
82+
}
83+
84+
// proj = Y_fixed^T * y_j (N_fixed × 1)
85+
Eigen::VectorXd proj = Y_fixed_map.transpose() * y_j;
86+
87+
// grad = GradientBasis * proj (num_dims × 1)
88+
Eigen::VectorXd grad = GradientBasis_map * proj;
89+
90+
// pinv_proj = pinvMat_fixed * proj (N_fixed × 1)
91+
Eigen::VectorXd pinv_proj = pinvMat_map * proj;
92+
93+
// Schur complement scalars
94+
double y_j_norm_sq = y_j.squaredNorm();
95+
double proj_pinv_proj = proj.dot(pinv_proj);
96+
97+
// s = ||y_j||^2/N + α - proj^T * pinvMat * proj / N^2
98+
double s = y_j_norm_sq / N + m_MinimumVariance - proj_pinv_proj / N_sq;
99+
100+
// Q[:, j] = (y_j - grad / N) / s
101+
for (int i = 0; i < num_dims; i++) {
102+
Q(i, j) = (y_j[i] - grad[i] / N) / s;
103+
}
104+
}
105+
}
106+
107+
// Now compute the updates from Q (same as original code)
108+
{
109+
vnl_matrix_type Jmatrix;
110+
vnl_matrix_type v;
111+
for (int j = 0; j < num_samples; j++) {
112+
int num = 0;
113+
int num2 = 0;
114+
for (unsigned int d = 0; d < m_DomainsPerShape; d++) {
115+
int dom = d + j * m_DomainsPerShape;
116+
if (c->GetDomainFlag(dom) == false) {
117+
if (d > 0) {
118+
num += m_AttributesPerDomain[d - 1] * c->GetNumberOfParticles(d - 1);
119+
if (m_UseXYZ[d - 1]) num += 3 * c->GetNumberOfParticles(d - 1);
120+
if (m_UseNormals[d - 1]) num += 3 * c->GetNumberOfParticles(d - 1);
121+
num2 += VDimension * c->GetNumberOfParticles(d - 1);
122+
}
123+
124+
int num_attr = m_AttributesPerDomain[d];
125+
if (m_UseXYZ[d]) num_attr += 3;
126+
if (m_UseNormals[d]) num_attr += 3;
127+
128+
Jmatrix.clear();
129+
Jmatrix.set_size(num_attr, VDimension);
130+
Jmatrix.fill(0.0);
131+
v.clear();
132+
v.set_size(num_attr, 1);
133+
v.fill(0.0);
134+
135+
for (unsigned int p = 0; p < c->GetNumberOfParticles(dom); p++) {
136+
v = Q.extract(num_attr, 1, num + p * num_attr, j);
137+
Jmatrix = m_ShapeGradient->extract(num_attr, 3, num + p * num_attr, 3 * j);
138+
vnl_matrix_type dx = Jmatrix.transpose() * v;
139+
for (unsigned int vd = 0; vd < VDimension; vd++) {
140+
m_PointsUpdate->put(num2 + p * VDimension + vd, j, dx(vd, 0));
141+
}
142+
}
143+
}
144+
}
145+
}
146+
}
147+
148+
// Compute energy using precomputed eigenvalues
149+
m_CurrentEnergy = 0.0;
150+
m_MinimumEigenValue = (*m_FixedW)(0, 0);
151+
for (int i = 0; i < m_PrecomputedNumFixedSamples; i++) {
152+
double w_i = (*m_FixedW)(i, i);
153+
double val = w_i * w_i / (m_PrecomputedNumFixedSamples - 1) + m_MinimumVariance;
154+
if (val < m_MinimumEigenValue) m_MinimumEigenValue = val;
155+
m_CurrentEnergy += log(val);
156+
}
157+
m_CurrentEnergy /= 2.0;
158+
159+
return; // Done with fixed shape space path
160+
}
161+
162+
// =========================================================================
163+
// ORIGINAL PATH: Full computation (no precomputed data or using mean energy)
164+
// =========================================================================
165+
// Use different profiling scopes to distinguish mean energy (init) vs entropy (optimize)
166+
TIME_SCOPE(m_UseMeanEnergy ? "ComputeUpdates::MeanEnergy" : "ComputeUpdates::Entropy");
167+
29168
vnl_matrix_type points_minus_mean(num_dims, num_samples, 0.0);
30169

31170
m_points_mean->clear();
@@ -44,7 +183,6 @@ void CorrespondenceFunction::ComputeUpdates(const ParticleSystem* c) {
44183

45184
if (this->m_UseMeanEnergy) {
46185
pinvMat.set_identity();
47-
m_InverseCovMatrix->setZero();
48186
} else {
49187
TIME_START("correspondence::gramMat");
50188

@@ -75,40 +213,9 @@ void CorrespondenceFunction::ComputeUpdates(const ParticleSystem* c) {
75213
pinvMat = (UG * invLambda) * UG.transpose();
76214
TIME_STOP("correspondence::pinvMat");
77215

78-
TIME_START("correspondence::lhs_rhs");
79-
80-
// Create Eigen maps for the VNL matrices
81-
Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>> UG_map(UG.data_block(),
82-
UG.rows(), UG.cols());
83-
// Create temporary Eigen matrices for the computations
84-
Eigen::MatrixXd projMat_eigen = points_minus_mean_map * UG_map;
85-
86-
// Convert invLambda diagonal matrix to Eigen
87-
// Since invLambda is a vnl_diag_matrix, we'll need to extract the diagonal
88-
Eigen::VectorXd invLambda_diag(invLambda.size());
89-
for (int i = 0; i < invLambda.size(); i++) {
90-
invLambda_diag(i) = invLambda(i, i);
91-
}
92-
Eigen::DiagonalMatrix<double, Eigen::Dynamic> invLambda_eigen(invLambda_diag);
93-
94-
// Perform the computations using Eigen
95-
const auto lhs = projMat_eigen * invLambda_eigen;
96-
const auto rhs = invLambda_eigen * projMat_eigen.transpose();
97-
TIME_STOP("correspondence::lhs_rhs");
98-
99-
// resize the inverse covariance matrix if necessary
100-
if (m_InverseCovMatrix->rows() != num_dims || m_InverseCovMatrix->cols() != num_dims) {
101-
m_InverseCovMatrix->resize(num_dims, num_dims);
102-
}
103-
104-
TIME_START("correspondence::covariance_multiply");
105-
if (PlatformUtils::is_macos()) {
106-
// this is about 2x faster on MacOS, but slower on Linux
107-
m_InverseCovMatrix->selfadjointView<Eigen::Upper>().rankUpdate(lhs);
108-
} else {
109-
m_InverseCovMatrix->noalias() = lhs * rhs;
110-
}
111-
TIME_STOP("correspondence::covariance_multiply");
216+
// Note: The full dM × dM inverse covariance matrix is NOT needed.
217+
// Per thesis equation 2.35, the gradient is simply: Y × (Y^T Y + αI)^{-1}
218+
// which is Y × pinvMat. No covariance_multiply required.
112219
}
113220

114221
vnl_matrix_type Q = points_minus_mean * pinvMat;
@@ -208,28 +315,15 @@ CorrespondenceFunction::VectorType CorrespondenceFunction::evaluate(unsigned int
208315
num += num1 * system->GetNumberOfParticles(i);
209316
}
210317

211-
vnl_matrix_type tmp1(sz_Yidx, sz_Yidx, 0.0);
212-
213-
if (this->m_UseMeanEnergy) {
214-
tmp1.set_identity();
215-
} else {
216-
// extract 3x3 submatrix at k,k
217-
Eigen::MatrixXd region = m_InverseCovMatrix->block(sz_Yidx, sz_Yidx, 3, 3);
218-
// convert to vnl
219-
for (unsigned int i = 0; i < 3; i++) {
220-
for (unsigned int j = 0; j < 3; j++) {
221-
tmp1(i, j) = region(i, j);
222-
}
223-
}
224-
}
225-
318+
// Per-particle energy: squared Euclidean distance from mean
319+
// This is used for line search step-size adaptation.
320+
// Note: We always use Euclidean distance (identity matrix) regardless of m_UseMeanEnergy,
321+
// because the full inverse covariance matrix is not computed.
226322
vnl_matrix_type Y_dom_idx(sz_Yidx, 1, 0.0);
227-
228323
Y_dom_idx = m_ShapeData->extract(sz_Yidx, 1, num, sampNum) - m_points_mean->extract(sz_Yidx, 1, num);
229324

230-
vnl_matrix_type tmp = Y_dom_idx.transpose() * tmp1;
231-
tmp *= Y_dom_idx;
232-
325+
// energy = ||y - mean||^2
326+
vnl_matrix_type tmp = Y_dom_idx.transpose() * Y_dom_idx;
233327
energy = tmp(0, 0);
234328

235329
maxdt = m_MinimumEigenValue;
@@ -246,4 +340,144 @@ CorrespondenceFunction::VectorType CorrespondenceFunction::evaluate(unsigned int
246340
return system->TransformVector(gradE, system->GetInversePrefixTransform(d) * system->GetInverseTransform(d));
247341
}
248342

343+
//---------------------------------------------------------------------------
344+
void CorrespondenceFunction::ClearPrecomputedFixedDomains() {
345+
m_HasPrecomputedFixedDomains = false;
346+
m_PrecomputedNumFixedSamples = 0;
347+
m_PrecomputedNumDims = 0;
348+
m_FixedMean.reset();
349+
m_FixedY.reset();
350+
m_FixedPinvMat.reset();
351+
m_FixedGradientBasis.reset();
352+
m_FixedU.reset();
353+
m_FixedW.reset();
354+
m_FixedDomainIndices.clear();
355+
}
356+
357+
//---------------------------------------------------------------------------
358+
void CorrespondenceFunction::PrecomputeForFixedDomains(const ParticleSystem* ps) {
359+
TIME_SCOPE("CorrespondenceFunction::PrecomputeForFixedDomains");
360+
361+
// Clear any existing precomputed data
362+
ClearPrecomputedFixedDomains();
363+
364+
// Count fixed and non-fixed shapes
365+
int total_shapes = m_ShapeData->cols();
366+
int num_fixed = 0;
367+
int num_unfixed = 0;
368+
369+
SW_LOG("PrecomputeForFixedDomains: ShapeData rows={}, cols={}, DomainsPerShape={}, PS domains={}",
370+
m_ShapeData->rows(), total_shapes, m_DomainsPerShape, ps->GetNumberOfDomains());
371+
372+
// Identify which shapes are fixed
373+
// A shape is considered fixed if ALL its domains are fixed
374+
for (int shape = 0; shape < total_shapes; shape++) {
375+
bool shape_is_fixed = true;
376+
for (int d = 0; d < m_DomainsPerShape; d++) {
377+
int dom = shape * m_DomainsPerShape + d;
378+
if (!ps->GetDomainFlag(dom)) {
379+
shape_is_fixed = false;
380+
break;
381+
}
382+
}
383+
if (shape_is_fixed) {
384+
m_FixedDomainIndices.push_back(shape);
385+
num_fixed++;
386+
} else {
387+
num_unfixed++;
388+
}
389+
}
390+
391+
// If no fixed shapes or all shapes are fixed, precomputation won't help
392+
if (num_fixed == 0 || num_unfixed == 0) {
393+
SW_LOG("PrecomputeForFixedDomains: No benefit from precomputation "
394+
"(fixed={}, unfixed={})", num_fixed, num_unfixed);
395+
return;
396+
}
397+
398+
SW_LOG("PrecomputeForFixedDomains: Precomputing for {} fixed shapes, {} unfixed shapes",
399+
num_fixed, num_unfixed);
400+
401+
int num_dims = m_ShapeData->rows();
402+
int num_fixed_samples = num_fixed;
403+
404+
// Store dimensions for validation later
405+
m_PrecomputedNumFixedSamples = num_fixed_samples;
406+
m_PrecomputedNumDims = num_dims;
407+
408+
// Step 1: Extract fixed shape data and compute mean
409+
// Build Y_fixed matrix: dM × N_fixed
410+
m_FixedY = std::make_shared<vnl_matrix_type>(num_dims, num_fixed_samples);
411+
m_FixedMean = std::make_shared<vnl_matrix_type>(num_dims, 1);
412+
m_FixedMean->fill(0.0);
413+
414+
// Copy fixed shape data
415+
for (int i = 0; i < num_fixed_samples; i++) {
416+
int fixed_shape_idx = m_FixedDomainIndices[i];
417+
for (int r = 0; r < num_dims; r++) {
418+
double val = m_ShapeData->get(r, fixed_shape_idx);
419+
m_FixedY->put(r, i, val);
420+
(*m_FixedMean)(r, 0) += val;
421+
}
422+
}
423+
424+
// Compute mean
425+
for (int r = 0; r < num_dims; r++) {
426+
(*m_FixedMean)(r, 0) /= num_fixed_samples;
427+
}
428+
429+
// Center the data: Y_fixed = Y_fixed - mean
430+
for (int i = 0; i < num_fixed_samples; i++) {
431+
for (int r = 0; r < num_dims; r++) {
432+
(*m_FixedY)(r, i) -= (*m_FixedMean)(r, 0);
433+
}
434+
}
435+
436+
// Step 2: Compute Gram matrix G = Y_fixed^T * Y_fixed
437+
TIME_START("precompute::gramMat");
438+
Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>
439+
Y_fixed_map(m_FixedY->data_block(), num_dims, num_fixed_samples);
440+
441+
Eigen::MatrixXd gramMat_eigen = Y_fixed_map.transpose() * Y_fixed_map;
442+
TIME_STOP("precompute::gramMat");
443+
444+
// Step 3: SVD of Gram matrix
445+
TIME_START("precompute::svd");
446+
vnl_matrix_type gramMat(num_fixed_samples, num_fixed_samples);
447+
for (int i = 0; i < num_fixed_samples; i++) {
448+
for (int j = 0; j < num_fixed_samples; j++) {
449+
gramMat(i, j) = gramMat_eigen(i, j);
450+
}
451+
}
452+
453+
vnl_svd<double> svd(gramMat);
454+
m_FixedU = std::make_shared<vnl_matrix_type>(svd.U());
455+
m_FixedW = std::make_shared<vnl_diag_matrix<double>>(svd.W());
456+
TIME_STOP("precompute::svd");
457+
458+
// Step 4: Compute regularized inverse: pinvMat = U * (W/(N-1) + αI)^{-1} * U^T
459+
TIME_START("precompute::pinvMat");
460+
vnl_diag_matrix<double> invLambda = *m_FixedW;
461+
for (int i = 0; i < num_fixed_samples; i++) {
462+
double val = invLambda(i, i) / (num_fixed_samples - 1) + m_MinimumVariance;
463+
invLambda(i, i) = 1.0 / val; // Invert
464+
}
465+
466+
m_FixedPinvMat = std::make_shared<vnl_matrix_type>(num_fixed_samples, num_fixed_samples);
467+
*m_FixedPinvMat = (*m_FixedU) * invLambda * m_FixedU->transpose();
468+
TIME_STOP("precompute::pinvMat");
469+
470+
// Step 5: Precompute gradient basis: M = Y_fixed * pinvMat
471+
// This is dM × N_fixed
472+
TIME_START("precompute::gradientBasis");
473+
m_FixedGradientBasis = std::make_shared<vnl_matrix_type>(num_dims, num_fixed_samples);
474+
*m_FixedGradientBasis = (*m_FixedY) * (*m_FixedPinvMat);
475+
TIME_STOP("precompute::gradientBasis");
476+
477+
m_HasPrecomputedFixedDomains = true;
478+
479+
SW_LOG("PrecomputeForFixedDomains: Precomputation complete. "
480+
"Fixed shapes: {}, Dimensions: {}", num_fixed_samples, num_dims);
481+
}
482+
249483
} // namespace shapeworks

0 commit comments

Comments
 (0)