Skip to content

Non-unique Games w/ Symmetry - #1391

Draft
j-c-c wants to merge 161 commits into
developfrom
nug
Draft

j-c-c wants to merge 161 commits into
developfrom
nug

Conversation

@j-c-c

@j-c-c j-c-c commented Jun 15, 2026 •

Copy link
Copy Markdown
Collaborator

Summary

 
Integrates the researcher-provided non-unique games (NUG) orientation-estimation code into ASPIRE-Python as CommonlineNUG. The estimator supports asymmetric (C1), cyclic (Cn), and dihedral (Dn) cases through ASPIRE’s source and OrientedSource workflows.
 

Changes beyond the researcher handoff

 

  • Replaced the standalone image-loading and projection workflow with ASPIRE sources, PolarFT, symmetry metadata, rotation conventions, and optional common-line shift search.
  • Removed dependencies on precomputed MATLAB files in the researcher’s data/ directory. The implementation now constructs the quaternion equality-constraint matrix and its inverse-related term, and computes the Fejer weight blocks from computed SO(3) grid and Wigner representations, rather than loading SO(3) grid, AEqJ.mat, and per-degree Fejer .mat files.
  • Made the implementation work through ASPIRE’s NumPy/CuPy backend selection, while moving small setup operations to the CPU where appropriate.
  • Improved coefficient computation by vectorizing the pairwise common-line loss; replaced dense permutation matrix operations with index based reordering; and replaced the large power-method calculation with an eigendecomposition of the smaller AI.T @ AI matrix.
  • Completed the end-to-end asymmetric (C1) pathway. The researcher’s handoff included a C1 ADMM solver; this integration adds C1 specific rotation recovery from its relaxed solution using ASPIRE’s synchronization-voting methods. It also processes Fejer inequality operations in image pair chunks to make larger (500 image) C1 GPU runs feasible.
  • Reduced C1 GPU memory pressure by processing Fejer inequality operations in image pair chunks and releasing unused CuPy cache before downstream reconstruction.
  • Reworked indexing, helper placement, preprocessing outputs, docstrings, logging, and progress reporting to make the port easier to read and maintain.
  • Added tests for C1/Cn/Dn orientation accuracy, dtype behavior, symmetry-aware rotation comparisons, block transformations, and the C1 batching boundary.
  • Added g_sync for symmetry-aware evaluation of estimated rotations. It retains the pre-existing cyclic-group path and adds finite-group synchronization for dihedral and other finite groups, with tests for C3, C4, D3, D4, T, and O.
  • Added galleries for simulated C1 projections of EMDB-2660 and the symmetric C4 and D3 examples, with the experimental class averages registered in ASPIRE’s downloader. These galleries are modeled off the experimental and simulated workflows described in the associated publications and produce comparable or better results.
     

Performance results

 
All experimental (and simulated) workflows from the associated publications are modeled within the uploaded gallery examples and have end-to-end runtimes 3-5x faster than the researcher provided code, with comparable or improved results with respect to rotation MSEs and reconstructed volume FSC scores.

@j-c-c j-c-c self-assigned this Jun 15, 2026
@codecov

codecov Bot commented Jun 15, 2026 •

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 97.22222% with 4 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.02%. Comparing base (e4a05cc) to head (eada3db).

Files with missing lines Patch % Lines
src/aspire/downloader/data_fetcher.py 33.33% 4 Missing ⚠️
Additional details and impacted files
@@             Coverage Diff             @@
##           develop    #1391      +/-   ##
===========================================
+ Coverage    90.65%   91.02%   +0.36%     
===========================================
  Files          135      136       +1     
  Lines        14932    16111    +1179     
===========================================
+ Hits         13537    14665    +1128     
- Misses        1395     1446      +51     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@j-c-c
j-c-c force-pushed the nug branch 2 times, most recently from 28ba09b to afd1d2b Compare June 23, 2026 15:29
@j-c-c
j-c-c force-pushed the nug branch 2 times, most recently from f96aa00 to 29ca799 Compare July 29, 2026 18:23
@j-c-c
j-c-c requested a review from garrettwrong July 29, 2026 19:59

@garrettwrong garrettwrong left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A few easy suggestions in-line but I do have larger concerns about the NUG code and am kicking it back so you can address them.

i) We don't appear to include a reproduction of original NUG for C1. If Amit is looking for complete coverage of groups' methods, this is required.

ii) We don't appear to include the symmetric extension authors reproducible example experiment. This is important since the method did not appear to work on the best experimental data we have available, so we will need this reproducible example as the full scientific test. The data will need to be made available for this purpose.

iii) After about line 300 or so into nug, where I oof'ed, the code becomes difficult to read/follow, with little documentation of the (many) variables. I had to start backing off the detailed review at that point, unfortunately. Not trying to reverse engineer that thing. While some docstrings are good, many are hit/miss, and missing coverage of their args/returns. Given (iv), this is sort of important. I didn't want to annotate every miss, too much.

iv) There are places where there are an excess of variables that could/should probably be restructured. This is normally the type of work the RSE's do. If we aren't restructuring the code to be improved (in some cases not merited/effective, I get it), I would expect the documentation to be better...

v) A lot of the xp code looks questionable. I have concerns that np was replaced with xp without considering the operations enough. Creating/performing many small array operations can often perform worse and is a poor use of limited resources. Even if performance is at par, if it fragments/uses GPU memory; this implementation may negatively impact later components that want to use the GPU for gains. I would want to see some evidence as to which of those areas actually effectively use the GPU and places that don't be reverted. Perhaps I'm totally misunderstanding some of the sizes/operations there... Some of the ones that look like might offer more were not converted, so I found it generally confusing.

vi) It is sort of confusing which methods were chosen to be static, which were not, and how they are nested (or not).

vii) a lot of the code has strange indexing. In some cases it would actually appear to be simpler as 0-based, but is written (like a port of) 1-based. I recognize a few of the areas might be cleaner as 1-based, and many neutral but are written 1-based. This is odd; reads like the code started in MATLAB and passed through multiple people who didn't clean up the indexing...

vii) I'm not sure how performance critical we need to be about this code, but the way it is written in some places has a lot of inefficient implementations. Lots of lists, scalars, lists of small arrays etc being constructed and passed around. Because of how the code is written, it isn't immediately clear where the bulk of compute is happening, so I can't offer much more there based on my limited review. (or if it matters...)

viii) The printing formatting, while fine, are not consistent with the rest of ASPIRE. In this file, there are multiple types of formatting. Probably better to convert to f-strings. FWIW, if it was not part of larger code base that uses the other formatting I would have no problem with the C style printf formats. Having multiple styles is not ideal.

I think there were some other things but I'm running out of steam/time. This is plenty. Some of those need to be addressed. Others, well, I'm not sure to what extent you've already improved from Ruiyi's code, or if the extra effort is worth it for this algorithm's results. We can discuss the right level of effort for this algo with Joakim and/or Amit when I get back, but you should have enough to stay busy until then.

Comment thread tests/test_nug.py Resolved
Comment thread src/aspire/abinitio/commonline_nug.py Resolved
Comment thread src/aspire/abinitio/commonline_nug.py Resolved
Comment thread src/aspire/abinitio/commonline_nug.py Resolved
Comment thread src/aspire/abinitio/commonline_nug.py Resolved
Comment thread src/aspire/abinitio/commonline_nug.py Resolved
Comment thread src/aspire/abinitio/commonline_nug.py Outdated

# Solve the symmetry-constrained SDP relaxation in the real, packed
# representation basis prepared by ADMM_preprocessing.
(

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

oof, what are these

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agreed this was difficult to follow. Preprocessing now returns a dictionary instead of a long tuple. Each group of related variables is unpacked with a description of its role in the solver.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, I'm afraid that shoving all the vars into a dictionary then immediately unpacking them is not really much of a code improvement. That's just the same thing with twice as much code 😂 . The explanations are helpful, thanks, but they look a bit much for code/line comments.

For these situations, it is probably better to create a simple class that contains the data and describe the variables there. Pass around and access instances of the data class.

Let me get through the rest of the August review to see if there are any other major items you'll need to address before I commit to fully reviewing the code.

Comment thread src/aspire/abinitio/commonline_nug.py Resolved
Comment thread src/aspire/abinitio/commonline_nug.py Resolved
Comment thread src/aspire/abinitio/commonline_nug.py Resolved
@garrettwrong

Copy link
Copy Markdown
Collaborator

Ampere CI should be restored now if you want to push any updates for this.

@j-c-c

j-c-c commented Oct 28, 2026

Copy link
Copy Markdown
Collaborator Author

@garrettwrong, Thanks for the initial review. The branch has changed a bit since, but here's how I addressed the larger concerns you raised:

i) C1 reproduction: The C1 workflow was incomplete from the researcher's provided code, which did contain 2 of the 3 main algorithm steps, coefficient computation (same as for symmetric case) and ADMM solver, but was missing the rotation recovery step. I integrated the researchers asymmetric ADMM solver and applied the same performance and code cleanup refactors that were applied to the symmetric solver. I then ported the rotation recovery steps from MATLAB. After discussion with Amit, it was determined that the MATLAB rotation recovery was implemented incorrectly, so I developed a new method in line with Amit's guidance. I added testing for the C1 code path and supporting methods. I also added a gallery example modeled off the simulated data experiments in the original C1 NUG paper (there were no real data experiments).

ii) Symmetric reproducible examples: Added C4 and D3 galleries using the authors’ 17 and 93 class averages, respectively. The datasets are available through ASPIRE’s downloader with registered URLs and hashes. The galleries estimate orientations and reconstruct volumes. I also added finite-group g_sync which allows for direct comparison of symmetry-synchronized rotations (the author used MSE of pairwise relative rotations, which are symmetry invariant).

iii - iv) Readability, documentation, and excess variables: Reorganized the code into labeled sections partitioning the code into main algorithm steps and helper sections. Also organized the ADMM methods into named sections to help with readability. Documented the (many) preprocessing values and return them in a dictionary rather than a long tuple. Expanded helper docstrings. Some ADMM update helpers remain nested because they operate on solver-local state. All methods have docstrings with param descriptions, with the exception of the nested functions whose docstrings contain only a method description.

v) Use of GPU resources: Audited small operations and moved suitable index, permutation, rank, and Wigner/Fejer setup work to NumPy while retaining CuPy for the larger ADMM operations. C1 Fejer operations now process image pairs in chunks, and unused CuPy cache is released after orientation estimation so downstream reconstruction can use the GPU. A 500-image C1 orientation-to-reconstruction handoff completed on our dev machines.

vi) Static, instance, and nested methods: Made helpers private where appropriate, used static methods for operations independent of estimator state, and moved rotation-recovery helpers out of their enclosing methods. ADMM update functions remain local to their respective solvers.

vii) Indexing: Went through and overhauled the indexing. Some of the 1-based indexing was from ported MATLAB code and some was tracking the mathematical Wigner degree which is 1-based. To make things clearer I used k_idx to represent Python zero-based indexing and k to represent the Wigner degree. Found and resolved a small indexing bug in the process.

vii.b) Performance inefficiencies: I profiled the main stages to identify where time and memory were being spent, then focused on those costs. I vectorized the pairwise loss calculation, replaced dense permutation matrix multiplications with index reordering, and replaced the memory heavy eigenvalue power method with an eigendecomposition of the smaller matrix. I haven’t removed every small allocation, but these changes address the major costs identified by profiling.

viii) Formatting: Using f-strings now throughout.

@j-c-c
j-c-c requested a review from garrettwrong October 28, 2026 19:11

@garrettwrong garrettwrong left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I got a little into it, but I think think you should address the return value situation again as I noted before I go through in detail and checkout/run things.

Also point me to the relevant lines/method names relating to this statement Some ADMM update helpers remain nested because they operate on solver-local state.. I have concerns about that (hopefully nothing-burgers 😇 ), which I'd like to suss out now, but don't want to chew through a thousand lines of code just to figure out where to start parsing. Thanks.

@j-c-c

j-c-c commented Oct 30, 2026

Copy link
Copy Markdown
Collaborator Author

I got a little into it, but I think think you should address the return value situation again as I noted before I go through in detail and checkout/run things.

Thanks, I'll get working on that.

Also point me to the relevant lines/method names relating to this statement Some ADMM update helpers remain nested because they operate on solver-local state.. I have concerns about that (hopefully nothing-burgers 😇 ), which I'd like to suss out now, but don't want to chew through a thousand lines of code just to figure out where to start parsing. Thanks.

The nested ADMM helpers are on lines 430–760 in the C1 solver and 927–1287 in the symmetric solver. They're in the sections labeled "Constraint operators", "ADMM update rules", and "Optional iteration logging".

I kept them local because they use constraints and index maps prepared by that solver. The changing ADMM variables are generally passed as arguments.

@garrettwrong

Copy link
Copy Markdown
Collaborator

Okay, yeah, generally I don't think there is good reason to define these as local functions with mixed scopes. I only looked at the one class, but saw enough of what I was concerned about right away. Perhaps they can be methods that belong to a data class holding that state instead of nested local functions?... or perhaps methods of the cl class that take a data class instance if needed? Some might make more sense one way or another. Some maybe don't need either (and some seemingly don't need to exist at all )....

Possibly some truly might benefit from this functional approach, and that's totally fine 😇 ; but its not my work to go through them one by one in a code review and sort that out for you.

With static class variables, data variables, and arguments (which are also seemingly changing with similar frequency to the data variables?) there can be at least 3 scopes in some of these functions, which is super confusing. I'm actually surprised that the linter did not reject. I have seen it reject similar style functional code, requiring use of kwargs. I would not want to have to test/debug them; potentially via live breakpoints only, yikes.

The first couple I took a look at have some unnecessary stuff... here's a quick one that I hope makes my point... I don't see why this needs to be a local function:

# in admm_plain_J()
       N=self.n_img  # don't think this changes...
       ....
       n_pairs = N * (N - 1) // 2  # so neither does this
       ...
        def pair_chunks():  # or this
            """
            Yield off-diagonal image-pair slices that bound GPU temporaries.
            """
            chunk_size = 1024
            for start in range(0, n_pairs, chunk_size):
                yield slice(start, min(start + chunk_size, n_pairs))

Moreover, unless I am missing something, this function doesn't even need to exist at all. A method that does this already exists in the standard python library as itertools.batched(some_iterable, chunk_size) or similar. The code that should exist instead ( chunk_size or batch_size being configured in the class init so the source code doesn't require editing...) is not implemented ...

While I appreciate the quantity of work here, I was hoping more of this sort of stuff would be unwound/refactored/addressed before investing your time in the testing/validation and my time by sending to me for review again ...

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants