Rough in code to write out NN graph as Weighted Adjacency List

Author: garrettwrong

src/aspire/classification/rir_class2d.py

@@ -29,6 +29,7 @@ def __init__(
         bispectrum_freq_cutoff=None,
         large_pca_implementation="legacy",
         nn_implementation="legacy",
+        output_nn_filename=None,
         bispectrum_implementation="legacy",
         aligner=None,
         dtype=None,
@@ -47,7 +48,8 @@ def __init__(
         Z. Zhao, Y. Shkolnisky, A. Singer, Rotationally Invariant Image Representation
         for Viewing Direction Classification in Cryo-EM. (2014)
 
-        :param src: Source instance
+        :param src: Source instance.  Note it is possible to use one `source` for classification (ie CWF),
+        and a different `source` for stacking in the `aligner`.
         :param pca_basis: Optional FSPCA Basis instance
         :param fspca_components: Components (top eigvals) to keep from full FSCPA, default truncates to  400.
         :param alpha: Amplitude Power Scale, default 1/3 (eq 20 from  RIIR paper).
@@ -119,6 +121,7 @@ def __init__(
                 f"Provided nn_implementation={nn_implementation} not in {nn_implementations.keys()}"
             )
         self._nn_classification = nn_implementations[nn_implementation]
+        self.output_nn_filename = output_nn_filename
 
         # # Do we have a sane Large Dataset PCA
         large_pca_implementations = {
@@ -185,6 +188,8 @@ def classify(self, diagnostics=False):
         # # Stage 2: Compute Nearest Neighbors
         logger.info("Calculate Nearest Neighbors")
         classes, reflections, distances = self.nn_classification(coef_b, coef_b_r)
+        if self.output_nn_filename is not None:
+            self._save_nn(classes, reflections, distances)
 
         if diagnostics:
             # Lets peek at the distribution of distances
@@ -351,7 +356,7 @@ def _legacy_nn_classification(self, coeff_b, coeff_b_r, batch_size=2000):
             # Check with Joakim about preference.
             # I (GBW) think class[i] should have class[i][0] be the original image index.
             classes[start:finish] = np.argsort(-corr, axis=1)[:, :n_nbor]
-            # Store the corr values for the n_nhors in this batch
+            # Store the corr values for the n_nbors in this batch
             distances[start:finish] = np.take_along_axis(
                 corr, classes[start:finish], axis=1
             )
@@ -366,6 +371,59 @@ def _legacy_nn_classification(self, coeff_b, coeff_b_r, batch_size=2000):
 
         return classes, refl, distances
 
+    def _save_nn(self, classes, reflections, distances):
+        """
+        Output the Nearest Neighbors graph as a weighted adjacency list.
+
+        Vertices are indexed by their natural index in `source`.
+        Note reflected images are represented by `index + src.n`.
+
+        Only the output of the Nearest Neighbor call is saved.
+        If you want a complete graph, specify 2*src.n neighbors,
+        that is all images and their reflections.
+
+        Because this is mixed datatypes (int and floating),
+        this will be output as a space delimited text file.
+
+        Vi1 Vj1 W_i1_j1 Vj2 Wi1_j2 ...
+        Vi2 Vj1 W_i2_j1 Vj2 Wi2_j2 ...
+        ...
+
+        """
+
+        # Construct the weighted adjacency list
+        AdjList = []
+        for k in range(len(classes)):
+
+            row = []
+            vik = classes[k][0]
+            row.append(vik)
+
+            for j in range(1, len(classes[k])):
+
+                # Neighbor index
+                vj = classes[k][j]
+                if reflections[k][j]:
+                    vj += self.src.n
+                row.append(vj)
+
+                # Neighbor Weight (distance)
+                wt = distances[k][j]
+                row.append(wt)
+
+            # Store this row of the AdjList
+            AdjList.append(row)
+
+        logger.info(
+            "Writing Nearest Neighbors as Weighted Adjacency List"
+            f" to {self.output_nn_filename}"
+        )
+
+        # Output
+        with open(self.output_nn_filename, "w") as fh:
+            for row in AdjList:
+                fh.write(" ".join(str(x) for x in row) + "\n")
+
     def _legacy_pca(self, M):
         """
         This is more or less the historic implementation ported