Implement O3 autofocus on light-sheet arm (#48)

* adding pythonnet and drift correction

* removing the pythonnet dependency that should live in copylot

* cleanup the demo code

* move drift_correction.py to examples

* add copylot to dependencies

* add O3 stage setup

* move serial numbers to acq_engine

* rename drift_correction script

* add defocus stack acquisition

* style

* update examples

* refactor z stage move

* add galvo scanning to acquire_ls_defocus_stack

* Update acquire_defocus_stack.py

* update examples

* switch to pylablib stage and relative moves

* add some logging

* add KIM101 compensation

* fix galvo reset position bug

* add find_focus.py example

* style

* ignore tests in examples dir

* move o3 refocus to acq_engine

* add checks and logging

* debug acq engine o3 refocus

* update example

* Create separate logs directory

* move conda env logger to logs

* rename to mantis_acquisition_log

* save acquired stacks

* implement timed o3 refocus

* update kim101 calibration

* add relative O3 travel limits

* logger fixes

* style

* add waveorder to deps and format pyproject

* update kim101 compensation factor

* add threshold and plotting to focus finding

* update waveorder dependency

* add microscope_operations documentation

* update data structure specs

* make acquire_ls_defocus_stack MantisAcquisition static method

* update examples

* rename_kim101 example

* rename z_range vars to avoid confusion

---------

Co-authored-by: Uditha Velidandla <[email protected]>

docs/data_structure.md

@@ -11,8 +11,6 @@ Organization of the raw data is constrained by the `pycromanager`-based acquisit
 
 YYYY_MM_DD <experiment_description>
 |--- <acq-name>_<n>
-|    |--- mantis_acquisition_log_YYYYMMDDTHHMMSS.txt
-|
 |    |--- positions.csv
 |
 |    |--- platemap.csv
@@ -31,6 +29,10 @@ YYYY_MM_DD <experiment_description>
 |        |--- <acq-name>_lightsheet_NDTiffStack_1.tif
 |        ...
 |
+|    |--- logs  # contains acquisition logs
+|        |--- mantis_acquisition_log_YYYYMMDDTHHMMSS.txt
+|        |--- conda_environment_log_YYYYMMDDTHHMMSS.txt
+|
 |--- <acq-name>_<n>  # one experiment folder may contain multiple acquisitions
 |    ...
 |
@@ -41,7 +43,7 @@ YYYY_MM_DD <experiment_description>
 
 ```
 
-An example dataset is provided in: `//ESS/comp_micro/rawdata/mantis/2023_02_21_mantis_dataset_standard/`.
+An example dataset is provided in: `//ESS/comp_micro/rawdata/mantis/2023_02_21_mantis_dataset_standard/`. (TODO: this example is now outdates)
 
 Each acquisition will contain a PTCZYX dataset; some dimensions may be singleton.
 

examples/acquire_defocus_stack.py

@@ -0,0 +1,48 @@
+import numpy as np
+from pycromanager import Core, Studio
+from mantis.acquisition.microscope_operations import (
+    setup_kim101_stage, 
+    acquire_ls_defocus_stack_and_display,
+    set_relative_kim101_position,
+)
+from waveorder.focus import focus_from_transverse_band
+
+mmc = Core()
+mmStudio = Studio()
+z_start = -105
+z_end = 105
+z_step = 15
+galvo = 'AP Galvo'
+galvo_range = [-0.5, 0, 0.5]
+
+z_stage = setup_kim101_stage('74000291')
+z_range = np.arange(z_start, z_end + z_step, z_step)
+
+# run 5 times over
+for i in range(5):
+    data = acquire_ls_defocus_stack_and_display(
+        mmc, 
+        mmStudio, 
+        z_stage, 
+        z_range,
+        galvo,
+        galvo_range,
+        close_display=False,
+    )
+
+    focus_indices = []
+    for stack in data:
+        idx = focus_from_transverse_band(
+            stack, NA_det=1.35, lambda_ill=0.55, pixel_size=6.5/40/1.4
+        )
+        focus_indices.append(idx)
+
+    valid_focus_indices = [idx for idx in focus_indices if idx is not None]
+    print(f'Valid focus indices: {valid_focus_indices}')
+
+    focus_idx = int(np.median(valid_focus_indices))
+    o3_displacement = int(z_range[focus_idx])
+    print(f'O3 displacement: {o3_displacement} steps')
+
+    set_relative_kim101_position(z_stage, o3_displacement)
+

examples/calibrate_kim101.py

@@ -0,0 +1,71 @@
+# Calibration procedure
+
+# Image a 1 um fluorescent bead with epi illumination and LS detection. Focus O3
+# on the bead. This script will defocus on one side of the bead and measure the
+# image intensity. The stage calibration factor is determined from the
+# difference in slope of average image intensity vs z position when traveling
+# in the positive or negative direction
+
+# This calibration procedure works alright, but could be improved
+
+#%%
+import numpy as np
+from pycromanager import Core, Studio
+import matplotlib.pyplot as plt
+from mantis.acquisition.microscope_operations import setup_kim101_stage, acquire_ls_defocus_stack_and_display
+
+#%%
+mmc = Core()
+mmStudio = Studio()
+z_stage = setup_kim101_stage('74000291')
+
+z_start = 0
+z_end = 105
+z_step = 15
+galvo = 'AP Galvo'
+galvo_range = [0]*5
+
+z_range = np.hstack(
+    (
+        np.arange(z_start, z_end + z_step, z_step),
+        np.arange(z_end, z_start - z_step, -z_step)
+    )
+)
+
+#%%
+data = acquire_ls_defocus_stack_and_display(
+    mmc, 
+    mmStudio, 
+    z_stage, 
+    z_range,
+    galvo,
+    galvo_range,
+    close_display = False
+)
+
+# %%
+steps_per_direction = len(z_range)//2
+intensity = data.max(axis=(-1, -2))
+
+pos_int = intensity[:, :steps_per_direction]
+pos_z = z_range[:steps_per_direction]
+
+neg_int = intensity[:, steps_per_direction:]
+neg_z = z_range[steps_per_direction:]
+
+A = np.vstack([pos_z, np.ones(len(pos_z))]).T
+pos_slope = []
+neg_slope = []
+for i in range(len(galvo_range)):
+    m, c = np.linalg.lstsq(A, pos_int[i], rcond=None)[0]
+    pos_slope.append(m)
+    m, c = np.linalg.lstsq(np.flipud(A), neg_int[i], rcond=None)[0]
+    neg_slope.append(m)
+
+compensation_factor = np.mean(pos_slope) / np.mean(neg_slope)
+print(compensation_factor)
+
+# %%
+plt.plot(intensity.flatten())
+
+# %%

examples/test_find_focus.py

@@ -0,0 +1,24 @@
+import os
+import tifffile
+import glob
+import napari
+import numpy as np
+from waveorder.focus import focus_from_transverse_band
+
+data_path = r'D:\2023_07_07_O3_autofocus'
+dataset = 'kidney_rfp_fov0'
+
+viewer = napari.Viewer()
+files = glob.glob(os.path.join(data_path, dataset, '*.ome.tif'))
+
+data = []
+points = []
+for i, file in enumerate(files):
+    stack = tifffile.imread(file, is_ome=False)
+    focus_idx = focus_from_transverse_band(stack, NA_det=1.35, lambda_ill=0.55, pixel_size=6.5/(40*1.4))
+    data.append(stack)
+    points.append([i, focus_idx, 50, 50])
+
+viewer.add_image(np.asarray(data))
+viewer.add_points(np.asarray(points), size=20)
+napari.run()

examples/test_kim101_copylot.py

@@ -0,0 +1,34 @@
+# This script tests controlling the KIM101 O3 stage using copylot. Ivan found
+# that copylot control of the stage runs into errors after ~100 relative moves
+# of the stage. We currently control the stage with pylablib and have not run
+# into such problems
+
+from copylot.hardware.stages.thorlabs.KIM001 import KCube_PiezoInertia
+import time
+from copylot import logger
+# from waveorder.focus import focus_from_transverse_band
+
+
+def test_labelfree_stage():
+    with KCube_PiezoInertia(serial_number='74000565', simulator=False) as stage_LF:
+        print(f'LF current position {stage_LF.position}')
+        stage_LF.move_relative(10)
+
+
+def test_light_sheet_stage():
+    ### LIGHT SHEET STAGE
+    with KCube_PiezoInertia(serial_number='74000291', simulator=False) as stage_LS:
+
+        # Change the acceleration and step rate
+        stage_LS.step_rate = 500
+        stage_LS.step_acceleration = 1000
+
+        # Test relative movement
+        step_size = -10
+        for i in range(10):
+            stage_LS.move_relative(10)
+        # stage_LS.move_relative(-step_size)
+
+
+if __name__ == '__main__':
+    test_light_sheet_stage()

examples/test_kim101_pylablib.py

@@ -0,0 +1,15 @@
+from pylablib.devices import Thorlabs
+devices = Thorlabs.list_kinesis_devices()
+
+stage = Thorlabs.KinesisPiezoMotor('74000291')
+
+p = stage.get_position()
+for i in range(50):
+    # stage.move_to(p+25); stage.wait_move()
+    # stage.move_to(p-25); stage.wait_move()
+
+    # relative moves work better
+    stage.move_by(25); stage.wait_move()
+    stage.move_by(-25); stage.wait_move()
+
+print('done')

mantis/acquisition/AcquisitionSettings.py

@@ -80,3 +80,6 @@ class MicroscopeSettings:
     use_autofocus: bool = False
     autofocus_stage: Optional[str] = None
     autofocus_method: Optional[str] = None
+    use_o3_refocus: bool = False
+    o3_refocus_config: Optional[ConfigSettings] = None
+    o3_refocus_interval_min: Optional[int] = None