TEST: finished adding xrf_data_1 tester

Author: hruth

src/pyxalign/alignment/projection_matching.py

@@ -1,6 +1,5 @@
 from functools import partial
 from typing import Callable, Optional
-import matplotlib
 import numpy as np
 import cupy as cp
 import copy
@@ -20,15 +19,14 @@
 import pyxalign.image_processing as ip
 import pyxalign.api.maps as maps
 from pyxalign.api.enums import DeviceType, MemoryConfig
-from pyxalign.api.options.alignment import ProjectionMatchingOptions, ProjectionMatchingPlotOptions
+from pyxalign.api.options.alignment import ProjectionMatchingOptions
 import pyxalign.gpu_utils as gutils
 from pyxalign.api.types import ArrayType, r_type
 from pyxalign.gpu_wrapper import device_handling_wrapper
 from IPython.display import clear_output
 import matplotlib.pyplot as plt
 from tqdm import tqdm
 import astra
-import time
 
 
 class ProjectionMatchingAligner(Aligner):

src/pyxalign/data_structures/task.py

@@ -4,6 +4,7 @@
 from PyQt5.QtWidgets import QApplication
 
 from pyxalign import gpu_utils
+from pyxalign.api.options.alignment import ProjectionMatchingOptions
 from pyxalign.data_structures.projections import (
     ComplexProjections,
     PhaseProjections,
@@ -37,7 +38,7 @@ def __init__(
         self.complex_projections = complex_projections
         self.phase_projections = phase_projections
         self.pma_object: ProjectionMatchingAligner = None
-        self.pma_GUI_list: list[ProjectionMatchingViewer] = []
+        self.pma_gui_list: list[ProjectionMatchingViewer] = []
 
     def get_cross_correlation_shift(
         self,
@@ -68,46 +69,40 @@ def get_cross_correlation_shift(
         projections.plot_staged_shift("Cross-correlation Shift")
         print("Cross-correlation shift stored in shift_manager")
 
-    def get_projection_matching_shift(self, initial_shift: Optional[np.ndarray] = None):
+    def get_projection_matching_shift(self, initial_shift: Optional[np.ndarray] = None) -> np.ndarray:
+        # clear existing astra objects
         if self.pma_object is not None:
             if hasattr(self.pma_object, "aligned_projections"):
-                # Clear old astra objects
                 self.pma_object.aligned_projections.volume.clear_astra_objects()
-
+        
+        # reset timers
         clear_timer_globals()
-        # Initialize the projection-matching alignment object
-        self.pma_object = ProjectionMatchingAligner(
-            self.phase_projections, self.options.projection_matching
-        )
+
+        # close old gui windows
         if self.options.projection_matching.interactive_viewer.close_old_windows:
-            self.clear_pma_GUI_list()
-        try:
-            if self.pma_object.options.interactive_viewer.update.enabled:
-                # Run PMA algorithm
-                shift = self.pma_object.run_with_GUI(initial_shift=initial_shift)
-                # Store the QWidget in a list so the window remains open even if
-                # another PMA loop is started
-                self.pma_GUI_list += [self.pma_object.gui]  # uncomment later
-                # Close window
-                # self.pma_object.gui.close() # I think adding this helped, or removing the list helped.
-            else:
-                # Run PMA algorithm
-                shift = self.pma_object.run(initial_shift=initial_shift)
-        except (Exception, KeyboardInterrupt):
-            shift = self.pma_object.total_shift * self.pma_object.scale
-        finally:
-            # Store the result in the ShiftManager object
-            self.phase_projections.shift_manager.stage_shift(
-                shift=shift,
-                function_type=enums.ShiftType.FFT,
-                alignment_options=self.options.projection_matching,
-            )
-            print("Projection-matching shift stored in shift_manager")
+            self.clear_pma_gui_list()
+        else:
+            self.pma_gui_list += [self.pma_object.gui]
+
+        # run the pma algorithm
+        self.pma_object, shift = run_projection_matching(
+            self.phase_projections, initial_shift, self.options.projection_matching
+        )
 
-    def clear_pma_GUI_list(self):
-        for gui in self.pma_GUI_list:
+        # Store the result in the ShiftManager object
+        self.phase_projections.shift_manager.stage_shift(
+            shift=shift,
+            function_type=enums.ShiftType.FFT,
+            alignment_options=self.options.projection_matching,
+        )
+        print("Projection-matching shift stored in shift_manager")
+
+        return shift
+
+    def clear_pma_gui_list(self):
+        for gui in self.pma_gui_list:
             gui.close()
-        self.pma_GUI_list = []
+        self.pma_gui_list = []
 
     def get_complex_projection_masks(self, enable_plotting: bool = False):
         clear_timer_globals()
@@ -150,3 +145,23 @@ def launch_viewer(self):
         app = QApplication.instance() or QApplication([])
         self.gui = TaskViewer(self)
         self.gui.show()
+
+
+def run_projection_matching(
+    phase_projections: PhaseProjections,
+    initial_shift: np.ndarray,
+    projection_matching_options: ProjectionMatchingOptions,
+) -> tuple[ProjectionMatchingAligner, np.ndarray]:
+    # Initialize the projection-matching alignment object
+    pma_object = ProjectionMatchingAligner(phase_projections, projection_matching_options)
+    try:
+        if pma_object.options.interactive_viewer.update.enabled:
+            # Run PMA algorithm
+            shift = pma_object.run_with_GUI(initial_shift=initial_shift)
+        else:
+            # Run PMA algorithm
+            shift = pma_object.run(initial_shift=initial_shift)
+    except (Exception, KeyboardInterrupt):
+        shift = pma_object.total_shift * pma_object.scale
+    finally:
+        return pma_object, shift

src/pyxalign/data_structures/xrf_task.py

@@ -1,33 +1,36 @@
 from typing import Optional
 import numpy as np
+from pyxalign import LaminographyAlignmentTask
 from pyxalign.alignment.cross_correlation import CrossCorrelationAligner
 from pyxalign.api import enums
 from pyxalign.api.options.device import DeviceOptions
 from pyxalign.api.options.projections import ProjectionOptions
 from pyxalign.api.options.task import AlignmentTaskOptions
+from pyxalign.data_structures.task import run_projection_matching
 from pyxalign.data_structures.xrf_projections import XRFProjections
 from pyxalign.timing.timer_utils import clear_timer_globals
 from pyxalign.api.types import r_type
 
 
 class XRFTask:
-    _primary_channel: str
-    projections_dict: dict[str, XRFProjections] = {}
-
     def __init__(
         self,
         xrf_array_dict: dict[str, np.ndarray],
         angles: np.ndarray,
         scan_numbers: np.ndarray,
         projection_options: ProjectionOptions,
-        task_options: AlignmentTaskOptions,
+        alignment_options: AlignmentTaskOptions,
         primary_channel: str,
     ):
+        self.pma_object = None
+        self.pma_gui_list = []
+        self._primary_channel: str
+        self.projections_dict: dict[str, XRFProjections] = {}
         # self.angles = angles
         # self.scan_numbers = scan_numbers
         self.channels = xrf_array_dict.keys()
         self.projection_options = projection_options
-        self.task_options = task_options
+        self.alignment_options = alignment_options
         self._primary_channel = primary_channel
         self.create_xrf_projections_object(xrf_array_dict, angles, scan_numbers)
         self._center_of_rotation = self.projections_dict[self._primary_channel].center_of_rotation
@@ -55,6 +58,10 @@ def create_xrf_projections_object(
                 # transform_tracker=
             )
 
+    # @property
+    # def phase_projections(self):
+    #     return self.projections_dict[self._primary_channel]
+
     @property
     def angles(self):
         return self.projections_dict[self._primary_channel].angles
@@ -87,26 +94,66 @@ def center_of_rotation(self):
     @center_of_rotation.setter
     def center_of_rotation(self, center_of_rotation: np.ndarray):
         self._center_of_rotation = center_of_rotation
-        for channel, projections in self.projections_dict.items():
+        for _, projections in self.projections_dict.items():
             projections.center_of_rotation = self._center_of_rotation * 1
 
     def apply_staged_shift_to_all_channels(self, device_options: Optional[DeviceOptions] = None):
-        for channel, projections in self.projections_dict.items():
+        for _, projections in self.projections_dict.items():
             projections.apply_staged_shift(device_options)
 
     def drop_projections_from_all_channels(self, remove_idx: list[int]):
-        for channel, projections in self.projections_dict.items():
+        for _, projections in self.projections_dict.items():
             projections.drop_projections(remove_idx=remove_idx)
 
     def pin_all_arrays(self):
-        for channel, projections in self.projections_dict.items():
+        for _, projections in self.projections_dict.items():
             projections.pin_arrays()
 
+    def get_projection_matching_shift(
+        self, initial_shift: Optional[np.ndarray] = None
+    ) -> np.ndarray:
+        # clear existing astra objects
+        if self.pma_object is not None:
+            if hasattr(self.pma_object, "aligned_projections"):
+                self.pma_object.aligned_projections.volume.clear_astra_objects()
+
+        # reset timers
+        clear_timer_globals()
+
+        # close old gui windows
+        if self.alignment_options.projection_matching.interactive_viewer.close_old_windows:
+            self.clear_pma_gui_list()
+        else:
+            self.pma_gui_list += [self.pma_object.gui]
+
+        # run the pma algorithm
+        self.pma_object, shift = run_projection_matching(
+            self.projections_dict[self.primary_channel],
+            initial_shift,
+            self.alignment_options.projection_matching,
+        )
+
+        # Save the resulting alignment shift
+        for _, projections in self.projections_dict.items():
+            projections.shift_manager.stage_shift(
+                shift=shift,
+                function_type=enums.ShiftType.FFT,
+                alignment_options=self.alignment_options.projection_matching,
+            )
+        print("Projection-matching shifts stored in shift_manager")
+
+        return shift
+
+    def clear_pma_gui_list(self):
+        for gui in self.pma_gui_list:
+            gui.close()
+        self.pma_gui_list = []
+
     def get_cross_correlation_shift(self, illum_sum: np.ndarray = None):
         clear_timer_globals()
         self.cross_correlation_aligner = CrossCorrelationAligner(
             projections=self.projections_dict[self._primary_channel],
-            options=self.task_options.cross_correlation,
+            options=self.alignment_options.cross_correlation,
         )
         # Placeholder for actual illum_sum
         if illum_sum is None:
@@ -121,7 +168,7 @@ def get_cross_correlation_shift(self, illum_sum: np.ndarray = None):
             projections.shift_manager.stage_shift(
                 shift=shift,
                 function_type=enums.ShiftType.CIRC,
-                alignment_options=self.task_options.cross_correlation,
+                alignment_options=self.alignment_options.cross_correlation,
             )
         projections.plot_staged_shift("Cross-correlation Shift")
         print("Cross-correlation shift stored in shift_manager")

tests/full_tests/cSAXS_e18044_LamNI_201907.py

@@ -280,7 +280,6 @@ def update_pma_options(pma_options: opts.ProjectionMatchingOptions, scale: int):
         )
 
         # Shift the projections by the projection-matching alignment shift
-        print(multi_gpu_device_options)
         task.phase_projections.apply_staged_shift(multi_gpu_device_options)
 
         # Save the fully aligned task