O3 autofocus improvement (#165)

* add calibration script

* add backlash correction and fix exp time bug

* remove O3 travel limits

* remove backlash correction

* expand O3 calibration script

* retry O3 autofocus with extended range

* small tweaks to O3 calibration script

* increase O3 autofocus FWHM threshold

* adjust backlash and increase O3 extended range search

* improve O3 autoexposure with better extended range scanning

* Refactor acquire_ls_defocus_stack

* run autoexposure before O3 autofocus

* set laser power right after autoexposure

* style

* move `github-markdown.css` to `mantis.acquisition.templates`

* move out scripts

* skip sorting in scripts

* Add option to exclude wells from O3 refocus and adjust how acquisition rates are updates after autoexposure

* style

* remove unused scripts import

* turn off live mode, check for illumination.csv file, and tweak o3 refocus var names

* update waveorder requirements

* style

* fix failing estimate_stabilization test

mantis/acquisition/AcquisitionSettings.py

@@ -141,6 +141,7 @@ class MicroscopeSettings:
     use_o3_refocus: bool = False
     o3_refocus_config: Optional[ConfigSettings] = None
     o3_refocus_interval_min: Optional[int] = None
+    o3_refocus_skip_wells: List[str] = field(default_factory=list)
 
 
 @dataclass

mantis/acquisition/acq_engine.py

@@ -40,6 +40,7 @@
 from mantis.acquisition.hook_functions.post_hardware_hook_functions import (
     log_acquisition_start,
     update_ls_hardware,
+    update_laser_power,
 )
 from mantis.acquisition.hook_functions.post_camera_hook_functions import (
     start_daq_counters,
@@ -61,6 +62,7 @@
 LS_CHANGE_TIME = 200  # time needed to change LS filter wheel, in ms
 LS_KIM101_SN = 74000291
 LF_KIM101_SN = 74000565
+KIM101_BACKLASH = 0  # backlash correction distance, in steps
 VORTRAN_488_COM_PORT = 'COM6'
 VORTRAN_561_COM_PORT = 'COM13'
 VORTRAN_639_COM_PORT = 'COM12'
@@ -209,6 +211,12 @@ def setup(self):
         Apply acquisition settings as specified by the class properties
         """
         if self.enabled:
+            # Turn off Live mode
+            if self.mmStudio:
+                snap_live_manager = self.mmStudio.get_snap_live_manager()
+                if snap_live_manager.is_live_mode_on():
+                    snap_live_manager.set_live_mode_on(False)
+
             # Apply microscope config settings
             for settings in self.microscope_settings.config_group_settings:
                 microscope_operations.set_config(
@@ -641,6 +649,13 @@ def setup_autoexposure(self):
             )
             return
 
+        if self.ls_acq.autoexposure_settings.autoexposure_method == 'manual':
+            # Check that the 'illumination.csv' file exists
+            if not (self._root_dir / 'illumination.csv').exists():
+                raise FileNotFoundError(
+                    f'The illumination.csv file required for manual autoexposure was not found in {self._root_dir}'
+                )
+
         # initialize lasers
         for channel_idx, config_name in enumerate(self.ls_acq.channel_settings.channels):
             if self.ls_acq.channel_settings.use_autoexposure[channel_idx]:
@@ -723,15 +738,10 @@ def go_to_position(self, position_index: int):
                     self.lf_acq.microscope_settings.autofocus_stage,
                 )
 
-    @staticmethod
     def acquire_ls_defocus_stack(
-        mmc: Core,
-        z_stage,
+        self,
         z_range: Iterable,
-        galvo: str,
         galvo_range: Iterable,
-        config_group: str = None,
-        config_name: str = None,
     ):
         """Acquire defocus stacks at different galvo positions and return image data
 
@@ -752,12 +762,23 @@ def acquire_ls_defocus_stack(
 
         """
         data = []
+        mmc = self.ls_acq.mmc
+        config_group = self.ls_acq.microscope_settings.o3_refocus_config.config_group
+        config_name = self.ls_acq.microscope_settings.o3_refocus_config.config_name
+        config_idx = self.ls_acq.channel_settings.channels.index(config_name)
+        exposure_time = self.ls_acq.channel_settings.default_exposure_times_ms[config_idx]
+        z_stage = self.ls_acq.o3_stage
+        galvo = self.ls_acq.slice_settings.z_stage_name
 
         # Set config
         if config_name is not None:
             mmc.set_config(config_group, config_name)
             mmc.wait_for_config(config_group, config_name)
 
+        # Set exposure time
+        if exposure_time is not None:
+            mmc.set_exposure(exposure_time)
+
         # Open shutter
         auto_shutter_state, shutter_state = microscope_operations.get_shutter_state(mmc)
         microscope_operations.open_shutter(mmc)
@@ -777,7 +798,9 @@ def acquire_ls_defocus_stack(
             mmc.set_position(galvo, p0 + p)
 
             # acquire defocus stack
-            z_stack = microscope_operations.acquire_defocus_stack(mmc, z_stage, z_range)
+            z_stack = microscope_operations.acquire_defocus_stack(
+                mmc, z_stage, z_range, backlash_correction_distance=KIM101_BACKLASH
+            )
             data.append(z_stack)
 
         # Reset camera triggering
@@ -793,22 +816,31 @@ def acquire_ls_defocus_stack(
 
         return np.asarray(data)
 
-    def refocus_ls_path(self):
+    def refocus_ls_path(
+        self, scan_left: bool = False, scan_right: bool = False
+    ) -> tuple[bool, bool, bool]:
         logger.info('Running O3 refocus algorithm on light-sheet arm')
+        success = False
 
         # Define O3 z range
         # 1 step is approx 20 nm, 15 steps are 300 nm which is sub-Nyquist sampling
         # The stack starts away from O2 and moves closer
         o3_z_start = -165
         o3_z_end = 165
         o3_z_step = 15
+        if scan_left:
+            logger.info('O3 refocus will scan further to the left')
+            o3_z_start *= 2
+        if scan_right:
+            logger.info('O3 refocus will scan further to the right')
+            o3_z_end *= 2
         o3_z_range = np.arange(o3_z_start, o3_z_end + o3_z_step, o3_z_step)
 
         # Define relative travel limits, in steps
         o3_z_stage = self.ls_acq.o3_stage
         target_z_position = o3_z_stage.true_position + o3_z_range
-        max_z_position = 750  # O3 is allowed to travel ~15 um towards O2
-        min_z_position = -1500  # O3 is allowed to travel ~30 um away from O2
+        max_z_position = np.inf  # O3 is allowed to travel ~15 um towards O2
+        min_z_position = -np.inf  # O3 is allowed to travel ~30 um away from O2
         if np.any(target_z_position > max_z_position) or np.any(
             target_z_position < min_z_position
         ):
@@ -827,15 +859,13 @@ def refocus_ls_path(self):
 
         # Acquire defocus stacks at several galvo positions
         data = self.acquire_ls_defocus_stack(
-            mmc=self.ls_acq.mmc,
-            z_stage=o3_z_stage,
             z_range=o3_z_range,
-            galvo=self.ls_acq.slice_settings.z_stage_name,
             galvo_range=galvo_range,
-            config_group=self.ls_acq.microscope_settings.o3_refocus_config.config_group,
-            config_name=self.ls_acq.microscope_settings.o3_refocus_config.config_name,
         )
 
+        # Discount O3 backlash compensation from true position count
+        o3_z_stage.true_position -= KIM101_BACKLASH * len(galvo_range)
+
         # Save acquired stacks in logs
         timestamp = datetime.now().strftime("%Y%m%dT%H%M%S")
         tifffile.imwrite(
@@ -847,11 +877,12 @@ def refocus_ls_path(self):
         wavelength = 0.55  # in um, approx
         # works well to distinguish between noise and sample when using z_step = 15
         # the idea is that true features in the sample will come in focus slowly
-        threshold_FWHM = 3.0
+        threshold_FWHM = 4.5
 
         focus_indices = []
+        peak_indices = []
         for stack_idx, stack in enumerate(data):
-            idx = focus_from_transverse_band(
+            idx, stats = focus_from_transverse_band(
                 stack,
                 NA_det=NA_DETECTION,
                 lambda_ill=wavelength,
@@ -860,6 +891,7 @@ def refocus_ls_path(self):
                 plot_path=self._logs_dir / f'ls_refocus_plot_{timestamp}_Pos{stack_idx}.png',
             )
             focus_indices.append(idx)
+            peak_indices.append(stats['peak_index'])
         logger.debug(
             'Stacks at galvo positions %s are in focus at slice %s',
             np.round(galvo_range, 3),
@@ -877,10 +909,27 @@ def refocus_ls_path(self):
             microscope_operations.set_relative_kim101_position(
                 self.ls_acq.o3_stage, o3_displacement
             )
+            success = True
         else:
             logger.error(
                 'Could not determine the correct O3 in-focus position. O3 will not move'
             )
+            if not any((scan_left, scan_right)):
+                # Only do this if we are not already scanning at an extended range
+                peak_indices = np.asarray(peak_indices)
+                max_idx = len(o3_z_range) - 1
+                if all(peak_indices < 0.2 * max_idx):
+                    scan_left = True
+                    logger.info(
+                        'O3 autofocus will scan further to the left at the next iteration'
+                    )
+                if all(peak_indices > 0.8 * max_idx):
+                    scan_right = True
+                    logger.info(
+                        'O3 autofocus will scan further to the right at the next iteration'
+                    )
+
+        return success, scan_left, scan_right
 
     def run_autoexposure(
         self,
@@ -1067,19 +1116,6 @@ def acquire(self):
                         )
                         continue
 
-                # O3 refocus
-                # Failing to refocus O3 will not abort the acquisition at the current PT index
-                if self.ls_acq.microscope_settings.use_o3_refocus:
-                    current_time = time.time()
-                    # Always refocus at the start
-                    if (
-                        (t_idx == 0 and p_idx == 0)
-                        or current_time - ls_o3_refocus_time
-                        > self.ls_acq.microscope_settings.o3_refocus_interval_min * 60
-                    ):
-                        self.refocus_ls_path()
-                        ls_o3_refocus_time = current_time
-
                 # autoexposure
                 if well_id != previous_well_id:
                     globals.new_well = True
@@ -1089,17 +1125,50 @@ def acquire(self):
                             well_id=well_id,
                             method=self.ls_acq.autoexposure_settings.autoexposure_method,
                         )
+
+                        # needs to be set before calling update_laser_power
+                        globals.ls_laser_powers = (
+                            self.ls_acq.channel_settings.laser_powers_per_well[well_id]
+                        )
+                        # This is a bit of a hack, laser powers should be set in update_ls_hardware
+                        for c_idx in range(self.ls_acq.channel_settings.num_channels):
+                            update_laser_power(
+                                self.ls_acq.channel_settings.light_sources, c_idx
+                            )
+
                     # Acq rate needs to be updated even if autoexposure was not rerun in this well
                     # Only do that if we are using autoexposure?
                     self.update_ls_acquisition_rates(
                         self.ls_acq.channel_settings.exposure_times_per_well[well_id]
                     )
+                    # needs to be set after calling update_ls_acquisition_rates
                     globals.ls_slice_acquisition_rates = (
                         self.ls_acq.slice_settings.acquisition_rate
                     )
-                    globals.ls_laser_powers = (
-                        self.ls_acq.channel_settings.laser_powers_per_well[well_id]
-                    )
+
+                # O3 refocus
+                # Failing to refocus O3 will not abort the acquisition at the current PT index
+                if self.ls_acq.microscope_settings.use_o3_refocus:
+                    current_time = time.time()
+                    # Always refocus at the start
+                    if (
+                        (t_idx == 0 and p_idx == 0)
+                        or current_time - ls_o3_refocus_time
+                        > self.ls_acq.microscope_settings.o3_refocus_interval_min * 60
+                    ):
+                        # O3 refocus can be skipped for certain wells
+                        if well_id in self.ls_acq.microscope_settings.o3_refocus_skip_wells:
+                            logger.debug(
+                                f'O3 refocus is due, but will be skipped in well {well_id}.'
+                            )
+                        else:
+                            success, scan_left, scan_right = self.refocus_ls_path()
+                            # If autofocus fails, try again with extended range if we know which way to go
+                            if not success and any((scan_left, scan_right)):
+                                success, _, _ = self.refocus_ls_path(scan_left, scan_right)
+                            # If it failed again, retry at the next position
+                            if success:
+                                ls_o3_refocus_time = current_time
 
                 # update events dictionaries
                 lf_events = deepcopy(lf_cz_events)

mantis/acquisition/hook_functions/post_hardware_hook_functions.py

@@ -25,10 +25,7 @@ def log_acquisition_start(events):
     return events
 
 
-def update_daq_freq(z_ctr_task, channels: list, events):
-    _event = get_first_acquisition_event(events)
-
-    c_idx = channels.index(_event['axes']['channel'])
+def update_daq_freq(z_ctr_task, c_idx: int):
     if z_ctr_task.is_task_done():
         z_ctr_task.stop()  # Counter needs to be stopped first
     z_ctr = z_ctr_task.co_channels[0]
@@ -38,13 +35,8 @@ def update_daq_freq(z_ctr_task, channels: list, events):
     logger.debug(f'Updating {z_ctr.name} pulse frequency to {acq_rates[c_idx]}')
     z_ctr.co_pulse_freq = acq_rates[c_idx]
 
-    return events
-
-
-def update_laser_power(lasers, channels: list, events):
-    _event = get_first_acquisition_event(events)
 
-    c_idx = channels.index(_event['axes']['channel'])
+def update_laser_power(lasers, c_idx: int):
     laser = lasers[c_idx]  # will be None if this channel does not use autoexposure
 
     if laser and globals.new_well:
@@ -55,15 +47,17 @@ def update_laser_power(lasers, channels: list, events):
         # Note, setting laser power takes ~1 s which is slow
         laser.pulse_power = laser_power
 
-    return events
-
 
-def update_ls_hardware(z_ctr_task, lasers, channels, events):
+def update_ls_hardware(z_ctr_task, lasers, channels: list, events):
     if not events:
         logger.debug('Acquisition events are not valid.')
         return
 
-    events = update_daq_freq(z_ctr_task, channels, events)
-    events = update_laser_power(lasers, channels, events)
+    _event = get_first_acquisition_event(events)
+    c_idx = channels.index(_event['axes']['channel'])
+
+    update_daq_freq(z_ctr_task, c_idx)
+    # As a hack, setting laser power after call to `run_autoexposure`
+    # update_laser_power(lasers, c_idx)
 
     return events

mantis/acquisition/microscope_operations.py

@@ -17,7 +17,7 @@
 
 logger = logging.getLogger(__name__)
 
-KIM101_COMPENSATION_FACTOR = 1.035
+KIM101_COMPENSATION_FACTOR = 1.0
 
 
 def _try_mmc_call(mmc, mmc_call_name, *mmc_carr_args):
@@ -334,6 +334,7 @@ def acquire_defocus_stack(
     datastore=None,
     channel_ind: int = 0,
     position_ind: int = 0,
+    backlash_correction_distance: int = 0,
 ):
     """Snap image at every z position and put image in a Micro-manager datastore
 
@@ -351,6 +352,8 @@ def acquire_defocus_stack(
         Channel index of acquired images in the Micro-manager datastore, by default 0
     position_ind : int, optional
         Position index of acquired images in the Micro-manager datastore, by default 0
+    backlash_correction: int, optional
+        Distance to add to the homing move of the stage to correct for backlash, by default 0
 
     Returns
     -------
@@ -397,7 +400,7 @@ def acquire_defocus_stack(
             datastore.put_image(image)
 
     # reset z stage
-    move_z(-relative_z_steps.sum())
+    move_z(-relative_z_steps.sum() + backlash_correction_distance)
 
     return np.asarray(data)
 

mantis/analysis/analyze_psf.py

@@ -20,7 +20,7 @@
 from numpy.typing import ArrayLike
 from scipy.signal import peak_widths
 
-import mantis.acquisition.scripts
+import mantis.acquisition.templates
 
 
 def _make_plots(
@@ -141,7 +141,7 @@ def generate_report(
     df_gaussian_fit.to_csv(output_path / 'psf_gaussian_fit.csv', index=False)
     df_1d_peak_width.to_csv(output_path / 'psf_1d_peak_width.csv', index=False)
 
-    with pkg_resources.path(mantis.acquisition.scripts, 'github-markdown.css') as css_path:
+    with pkg_resources.path(mantis.acquisition.templates, 'github-markdown.css') as css_path:
         shutil.copy(css_path, output_path)
     html_file_path = output_path / ('psf_analysis_report.html')
     with open(html_file_path, 'w') as file:

mantis/cli/estimate_stabilization.py

@@ -47,7 +47,7 @@ def estimate_position_focus(
             if np.sum(data_zyx) == 0:
                 focal_plane = 0
             else:
-                focal_plane = focus_from_transverse_band(
+                focal_plane, _ = focus_from_transverse_band(
                     data_zyx,
                     NA_det=NA_DET,
                     lambda_ill=LAMBDA_ILL,