From f98ebd009cb879af69f6d6b435e00914c8c02166 Mon Sep 17 00:00:00 2001
From: Uditha Velidandla <edhiratam@gmail.com>
Date: Wed, 21 Jun 2023 16:16:35 -0700
Subject: [PATCH 01/45] adding pythonnet and drift correction

---
 mantis/analysis/drift_correction.py | 27 +++++++++++++++++++++++++++
 pyproject.toml                      |  1 +
 2 files changed, 28 insertions(+)
 create mode 100644 mantis/analysis/drift_correction.py

diff --git a/mantis/analysis/drift_correction.py b/mantis/analysis/drift_correction.py
new file mode 100644
index 00000000..72e86dcd
--- /dev/null
+++ b/mantis/analysis/drift_correction.py
@@ -0,0 +1,27 @@
+# %%
+from copylot.hardware.stages.thorlabs.KIM001 import KCube_PiezoInertia
+import time
+from copylot import logger
+
+### LABEL_FREE STAGE
+# stage_LF = KCube_PiezoInertia(serial_number='74000565', simulator=False)
+
+# Instantiate the piezo stages
+with KCube_PiezoInertia(serial_number='74000291', simulator=False) as stage_LS:
+    # Test Moving
+    # print(f'LF current position {stage_LF.position}')
+    # stage_LF.move_relative(10)
+    print(f'LS current position {stage_LS.position}')
+    stage_LS.move_relative(100)
+    stage_LS.move_relative(-100)
+    stage_LS.step_rate = 50
+    stage_LS.step_acceleration = 200
+    print(f'acceleration {stage_LS.step_acceleration} rate {stage_LS.step_rate}')
+    stage_LS.move_relative(100)
+    stage_LS.move_relative(-100)
+
+# stage_LS = KCube_PiezoInertia(serial_number='74000291', simulator=False)
+# print(f'LS current position {stage_LS.position}')
+# stage_LS.move_relative(100)
+# stage_LS.move_relative(-100)
+# stage_LS.disconnect()
\ No newline at end of file
diff --git a/pyproject.toml b/pyproject.toml
index 8ea6cf41..738b9380 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -31,6 +31,7 @@ dependencies = [
     "nidaqmx",
     "iohub==0.1.0.dev3",
     "matplotlib",
+    "pythonnet"
 ]
 
 

From 541b816f5e8de5eb9351594f896e0ade4f45d880 Mon Sep 17 00:00:00 2001
From: Eduardo Hirata-Miyasaki <edhiratam@gmail.com>
Date: Wed, 21 Jun 2023 16:31:35 -0700
Subject: [PATCH 02/45] removing the pythonnet dependency that should live in
 copylot

---
 mantis/analysis/drift_correction.py | 18 ++++++++----------
 pyproject.toml                      |  3 +--
 2 files changed, 9 insertions(+), 12 deletions(-)

diff --git a/mantis/analysis/drift_correction.py b/mantis/analysis/drift_correction.py
index 72e86dcd..0cdb57ce 100644
--- a/mantis/analysis/drift_correction.py
+++ b/mantis/analysis/drift_correction.py
@@ -4,24 +4,22 @@
 from copylot import logger
 
 ### LABEL_FREE STAGE
-# stage_LF = KCube_PiezoInertia(serial_number='74000565', simulator=False)
-
-# Instantiate the piezo stages
-with KCube_PiezoInertia(serial_number='74000291', simulator=False) as stage_LS:
-    # Test Moving
+# with KCube_PiezoInertia(serial_number='74000565', simulator=False) as stage_LF:
     # print(f'LF current position {stage_LF.position}')
     # stage_LF.move_relative(10)
+
+### LIGHT SHEET STAGE
+with KCube_PiezoInertia(serial_number='74000291', simulator=False) as stage_LS:
+    # Test the relative movement
     print(f'LS current position {stage_LS.position}')
     stage_LS.move_relative(100)
     stage_LS.move_relative(-100)
+
+    # Change the acceleration and step rate
     stage_LS.step_rate = 50
     stage_LS.step_acceleration = 200
     print(f'acceleration {stage_LS.step_acceleration} rate {stage_LS.step_rate}')
+    # Test the movement with the different acceleration
     stage_LS.move_relative(100)
     stage_LS.move_relative(-100)
 
-# stage_LS = KCube_PiezoInertia(serial_number='74000291', simulator=False)
-# print(f'LS current position {stage_LS.position}')
-# stage_LS.move_relative(100)
-# stage_LS.move_relative(-100)
-# stage_LS.disconnect()
\ No newline at end of file
diff --git a/pyproject.toml b/pyproject.toml
index 738b9380..16b19819 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -30,8 +30,7 @@ dependencies = [
     "pycromanager==0.25.40",
     "nidaqmx",
     "iohub==0.1.0.dev3",
-    "matplotlib",
-    "pythonnet"
+    "matplotlib"
 ]
 
 

From e75eb99f32ad76653a90348874e948f0ebadc973 Mon Sep 17 00:00:00 2001
From: Eduardo Hirata-Miyasaki <edhiratam@gmail.com>
Date: Wed, 21 Jun 2023 16:36:50 -0700
Subject: [PATCH 03/45] cleanup the demo code

---
 mantis/analysis/drift_correction.py | 41 +++++++++++++++++------------
 1 file changed, 24 insertions(+), 17 deletions(-)

diff --git a/mantis/analysis/drift_correction.py b/mantis/analysis/drift_correction.py
index 0cdb57ce..ced92ac3 100644
--- a/mantis/analysis/drift_correction.py
+++ b/mantis/analysis/drift_correction.py
@@ -2,24 +2,31 @@
 from copylot.hardware.stages.thorlabs.KIM001 import KCube_PiezoInertia
 import time
 from copylot import logger
+from waveorder.focus import focus_from_transverse_band
 
-### LABEL_FREE 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)
 
-### LIGHT SHEET STAGE
-with KCube_PiezoInertia(serial_number='74000291', simulator=False) as stage_LS:
-    # Test the relative movement
-    print(f'LS current position {stage_LS.position}')
-    stage_LS.move_relative(100)
-    stage_LS.move_relative(-100)
+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)
 
-    # Change the acceleration and step rate
-    stage_LS.step_rate = 50
-    stage_LS.step_acceleration = 200
-    print(f'acceleration {stage_LS.step_acceleration} rate {stage_LS.step_rate}')
-    # Test the movement with the different acceleration
-    stage_LS.move_relative(100)
-    stage_LS.move_relative(-100)
 
+def test_light_sheet_stage():
+    ### LIGHT SHEET STAGE
+    with KCube_PiezoInertia(serial_number='74000291', simulator=False) as stage_LS:
+        # Test the relative movement
+        print(f'LS current position {stage_LS.position}')
+        stage_LS.move_relative(100)
+        stage_LS.move_relative(-100)
+
+        # Change the acceleration and step rate
+        stage_LS.step_rate = 50
+        stage_LS.step_acceleration = 200
+        print(f'acceleration {stage_LS.step_acceleration} rate {stage_LS.step_rate}')
+        # Test the movement with the different acceleration
+        stage_LS.move_relative(100)
+        stage_LS.move_relative(-100)
+
+
+if __name__ == '__main__':
+    test_light_sheet_stage()

From 6e6b2f1055d41bcb5b8b4cc5d7504fbf26e13b22 Mon Sep 17 00:00:00 2001
From: Uditha Velidandla <ivan.ivanov@czbiohub.org>
Date: Thu, 29 Jun 2023 17:09:53 -0700
Subject: [PATCH 04/45] move drift_correction.py to examples

---
 {mantis/analysis => examples}/drift_correction.py | 0
 1 file changed, 0 insertions(+), 0 deletions(-)
 rename {mantis/analysis => examples}/drift_correction.py (100%)

diff --git a/mantis/analysis/drift_correction.py b/examples/drift_correction.py
similarity index 100%
rename from mantis/analysis/drift_correction.py
rename to examples/drift_correction.py

From 23d9344e6f73c89ade41e318a0a580aab62502b9 Mon Sep 17 00:00:00 2001
From: Uditha Velidandla <ivan.ivanov@czbiohub.org>
Date: Thu, 29 Jun 2023 17:27:15 -0700
Subject: [PATCH 05/45] add copylot to dependencies

---
 examples/drift_correction.py | 18 +++++++++---------
 pyproject.toml               |  3 ++-
 2 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/examples/drift_correction.py b/examples/drift_correction.py
index ced92ac3..24953afa 100644
--- a/examples/drift_correction.py
+++ b/examples/drift_correction.py
@@ -2,7 +2,7 @@
 from copylot.hardware.stages.thorlabs.KIM001 import KCube_PiezoInertia
 import time
 from copylot import logger
-from waveorder.focus import focus_from_transverse_band
+# from waveorder.focus import focus_from_transverse_band
 
 
 def test_labelfree_stage():
@@ -14,18 +14,18 @@ def test_labelfree_stage():
 def test_light_sheet_stage():
     ### LIGHT SHEET STAGE
     with KCube_PiezoInertia(serial_number='74000291', simulator=False) as stage_LS:
-        # Test the relative movement
+
         print(f'LS current position {stage_LS.position}')
-        stage_LS.move_relative(100)
-        stage_LS.move_relative(-100)
 
         # Change the acceleration and step rate
-        stage_LS.step_rate = 50
-        stage_LS.step_acceleration = 200
+        stage_LS.step_rate = 500
+        stage_LS.step_acceleration = 1000
         print(f'acceleration {stage_LS.step_acceleration} rate {stage_LS.step_rate}')
-        # Test the movement with the different acceleration
-        stage_LS.move_relative(100)
-        stage_LS.move_relative(-100)
+
+        # Test relative movement
+        step_size = 10
+        stage_LS.move_relative(step_size)
+        stage_LS.move_relative(-step_size)
 
 
 if __name__ == '__main__':
diff --git a/pyproject.toml b/pyproject.toml
index 16b19819..b164937e 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -30,7 +30,8 @@ dependencies = [
     "pycromanager==0.25.40",
     "nidaqmx",
     "iohub==0.1.0.dev3",
-    "matplotlib"
+    "coPylot@git+https://github.com/czbiohub-sf/coPylot.git#egg=thorlabs_stage",
+    "matplotlib",
 ]
 
 

From 2f83900c8150d58fd191a5e86f19e34df27bc755 Mon Sep 17 00:00:00 2001
From: Uditha Velidandla <ivan.ivanov@czbiohub.org>
Date: Thu, 29 Jun 2023 17:47:34 -0700
Subject: [PATCH 06/45] add O3 stage setup

---
 mantis/acquisition/microscope_operations.py | 23 +++++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/mantis/acquisition/microscope_operations.py b/mantis/acquisition/microscope_operations.py
index 840cd865..2af08d99 100644
--- a/mantis/acquisition/microscope_operations.py
+++ b/mantis/acquisition/microscope_operations.py
@@ -7,8 +7,13 @@
 
 from nidaqmx.constants import AcquisitionType
 
+from copylot.hardware.stages.thorlabs.KIM001 import KCube_PiezoInertia
+
 logger = logging.getLogger(__name__)
 
+LS_KIM101_SN = 74000291
+LF_KIM101_SN = 74000565
+
 
 def _try_mmc_call(mmc, mmc_call_name, *mmc_carr_args):
     """Wrapper that tries to repeat calls to mmCore if they fail. Largely copied
@@ -218,3 +223,21 @@ def autofocus(mmc, mmStudio, z_stage_name: str, z_position):
         logger.error(f'Autofocus call failed after {len(z_offsets)} tries')
 
     return autofocus_success
+
+def setup_ls_o3_stage(step_rate = 500, step_acceleration = 1000):
+    stage = KCube_PiezoInertia(serial_number=str(LS_KIM101_SN))
+
+    # Change the acceleration and step rate
+    stage.step_rate = step_rate
+    stage.step_acceleration = step_acceleration
+
+    return stage
+
+def setup_lf_o3_stage(step_rate = 500, step_acceleration = 1000):
+    stage = KCube_PiezoInertia(serial_number=str(LF_KIM101_SN))
+
+    # Change the acceleration and step rate
+    stage.step_rate = step_rate
+    stage.step_acceleration = step_acceleration
+
+    return stage

From 9d6a059a0f1f83f9ede33f8b18f2d364c10b1ad3 Mon Sep 17 00:00:00 2001
From: Uditha Velidandla <ivan.ivanov@czbiohub.org>
Date: Thu, 29 Jun 2023 17:50:44 -0700
Subject: [PATCH 07/45] move serial numbers to acq_engine

---
 mantis/acquisition/acq_engine.py            |  2 ++
 mantis/acquisition/microscope_operations.py | 16 ++--------------
 2 files changed, 4 insertions(+), 14 deletions(-)

diff --git a/mantis/acquisition/acq_engine.py b/mantis/acquisition/acq_engine.py
index 49adb8a8..0e47a9b3 100644
--- a/mantis/acquisition/acq_engine.py
+++ b/mantis/acquisition/acq_engine.py
@@ -47,6 +47,8 @@
 MCL_STEP_TIME = 1.5  # in ms
 LC_CHANGE_TIME = 20  # in ms
 LS_CHANGE_TIME = 200  # time needed to change LS filter wheel, in ms
+LS_KIM101_SN = 74000291
+LF_KIM101_SN = 74000565
 
 logger = logging.getLogger(__name__)
 
diff --git a/mantis/acquisition/microscope_operations.py b/mantis/acquisition/microscope_operations.py
index 2af08d99..1c85b240 100644
--- a/mantis/acquisition/microscope_operations.py
+++ b/mantis/acquisition/microscope_operations.py
@@ -11,9 +11,6 @@
 
 logger = logging.getLogger(__name__)
 
-LS_KIM101_SN = 74000291
-LF_KIM101_SN = 74000565
-
 
 def _try_mmc_call(mmc, mmc_call_name, *mmc_carr_args):
     """Wrapper that tries to repeat calls to mmCore if they fail. Largely copied
@@ -224,17 +221,8 @@ def autofocus(mmc, mmStudio, z_stage_name: str, z_position):
 
     return autofocus_success
 
-def setup_ls_o3_stage(step_rate = 500, step_acceleration = 1000):
-    stage = KCube_PiezoInertia(serial_number=str(LS_KIM101_SN))
-
-    # Change the acceleration and step rate
-    stage.step_rate = step_rate
-    stage.step_acceleration = step_acceleration
-
-    return stage
-
-def setup_lf_o3_stage(step_rate = 500, step_acceleration = 1000):
-    stage = KCube_PiezoInertia(serial_number=str(LF_KIM101_SN))
+def setup_kim101_stage(serial_number: int, step_rate = 500, step_acceleration = 1000):
+    stage = KCube_PiezoInertia(serial_number=str(serial_number))
 
     # Change the acceleration and step rate
     stage.step_rate = step_rate

From dd226dec8b1cfab1738046cd0768191ad39f0e7f Mon Sep 17 00:00:00 2001
From: Uditha Velidandla <ivan.ivanov@czbiohub.org>
Date: Fri, 30 Jun 2023 11:35:44 -0700
Subject: [PATCH 08/45] rename drift_correction script

---
 examples/{drift_correction.py => test_kim101_stage.py} | 0
 1 file changed, 0 insertions(+), 0 deletions(-)
 rename examples/{drift_correction.py => test_kim101_stage.py} (100%)

diff --git a/examples/drift_correction.py b/examples/test_kim101_stage.py
similarity index 100%
rename from examples/drift_correction.py
rename to examples/test_kim101_stage.py

From 7d3b0d0aee5d3e1476b162c26347d606a921d50d Mon Sep 17 00:00:00 2001
From: Uditha Velidandla <ivan.ivanov@czbiohub.org>
Date: Fri, 30 Jun 2023 13:24:27 -0700
Subject: [PATCH 09/45] add defocus stack acquisition

---
 examples/acquire_defocus_stack.py           |  13 ++
 mantis/acquisition/microscope_operations.py | 191 +++++++++++++++++++-
 2 files changed, 203 insertions(+), 1 deletion(-)
 create mode 100644 examples/acquire_defocus_stack.py

diff --git a/examples/acquire_defocus_stack.py b/examples/acquire_defocus_stack.py
new file mode 100644
index 00000000..807edd96
--- /dev/null
+++ b/examples/acquire_defocus_stack.py
@@ -0,0 +1,13 @@
+from pycromanager import Core, Studio
+from mantis.acquisition.microscope_operations import acquire_ls_defocus_stack
+
+mmc = Core()
+mmStudio = Studio()
+z_stage = 'Z'
+z_start = 0
+z_end = 10
+z_step = 1
+
+data = acquire_ls_defocus_stack(mmc, mmStudio, z_stage, z_start, z_end, z_step)
+
+print(data.shape)
diff --git a/mantis/acquisition/microscope_operations.py b/mantis/acquisition/microscope_operations.py
index 1c85b240..b9364d96 100644
--- a/mantis/acquisition/microscope_operations.py
+++ b/mantis/acquisition/microscope_operations.py
@@ -1,12 +1,15 @@
 import logging
 import time
+import numpy as np
+from pycromanager import Core, Studio
 
-from typing import Tuple
+from typing import Tuple, Iterable
 
 import nidaqmx
 
 from nidaqmx.constants import AcquisitionType
 
+from copylot.hardware.stages.abstract_stage import AbstractStage
 from copylot.hardware.stages.thorlabs.KIM001 import KCube_PiezoInertia
 
 logger = logging.getLogger(__name__)
@@ -229,3 +232,189 @@ def setup_kim101_stage(serial_number: int, step_rate = 500, step_acceleration =
     stage.step_acceleration = step_acceleration
 
     return stage
+
+def create_ram_datastore(
+    mmStudio: Studio,
+):
+    datastore = mmStudio.get_data_manager().create_ram_datastore()
+    mmStudio.get_display_manager().create_display(datastore)
+
+    return datastore
+
+def acquire_defocus_stack(
+    mmc: Core,
+    mmStudio: Studio,
+    datastore,
+    z_stage,
+    z_range: Iterable,
+    channel_ind: int=0,
+    position_ind: int=0,
+):
+    """Snap image at every z position and put image in the datastore
+
+    Parameters
+    ----------
+    mmc : Core
+    mmStudio : Studio
+    datastore : micromanager.data.Datastore
+        Micro-manager datastore object
+    z_stage : str or coPylot stage object
+    z_range : Iterable
+    channel_ind : int, optional
+        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
+
+    Returns
+    -------
+    data : np.array
+
+    """
+    data = []
+
+    for z_ind, z in enumerate(z_range):
+        # set z position
+        if isinstance(z_stage, str):
+            # this is a MM stage
+            mmc.set_relative_position(z_stage, float(z))
+        elif issubclass(z_stage, AbstractStage):
+            # this is a copylot stage
+            z_stage.move_relative(z)
+        else:
+            raise RuntimeError(f'Unknown z stage: {z_stage}')
+
+        # snap image
+        mmc.snap_image()
+        tagged_image = mmc.get_tagged_image()
+
+        # get image data
+        image_data = np.reshape(
+            tagged_image.pix, 
+            (tagged_image.tags['Height'], tagged_image.tags['Width'])
+        )
+        data.append(image_data.astype('uint16'))
+
+        # set image coordinates and put in datastore
+        image = mmStudio.get_data_manager().convert_tagged_image(tagged_image)
+        coords_builder = image.get_coords().copy()
+        coords_builder = coords_builder.z(z_ind)
+        coords_builder = coords_builder.channel(channel_ind)
+        coords_builder = coords_builder.stage_position(position_ind)
+        mm_coords = coords_builder.build()
+
+        image = image.copy_at_coords(mm_coords)
+        datastore.put_image(image)
+
+    return np.asarray(data)
+
+def acquire_ls_defocus_stack(
+    mmc: Core,
+    mmStudio: Studio,
+    z_stage,
+    z_start: float, 
+    z_end: float, 
+    z_step: float,
+    config_group: str=None, 
+    config_name: str=None,
+):
+    """Acquire defocus stacks at different galvo positions
+
+    Parameters
+    ----------
+    mmc : Core
+        _description_
+    mmStudio : Studio
+        _description_
+    z_stage : _type_
+        _description_
+    z_start : float
+        _description_
+    z_end : float
+        _description_
+    z_step : float
+        _description_
+    config_group : str, optional
+        _description_, by default None
+    config_name : str, optional
+        _description_, by default None
+
+    Returns
+    -------
+    data : np.array
+
+    """
+    datastore = create_ram_datastore(mmStudio)
+    z_range = np.arange(z_start, z_end+z_step, z_step)
+    data = []
+
+    # Set config
+    if config_name is not None:
+        mmc.set_config(config_group, config_name)
+        mmc.wait_for_config(config_group, config_name)
+
+    # Open shutter
+    auto_shutter_state, shutter_state = get_shutter_state(mmc)
+    open_shutter(mmc)
+
+    # acquire stack at different galvo positions
+    for p_idx, p in enumerate(range(3)):
+        # TODO: set galvo position
+
+        # acquire defocus stack
+        z_stack = acquire_defocus_stack(
+            mmc, mmStudio, datastore, z_stage, z_range, channel_ind=0, position_ind=p_idx
+        )
+        data.append(z_stack)
+
+    # freeze datastore to indicate that we are finished writing to it
+    datastore.freeze()
+
+    # Reset shutter
+    reset_shutter(mmc, auto_shutter_state, shutter_state)
+
+    return np.asarray(data)
+
+def get_shutter_state(mmc: Core):
+    """Return the current state of the shutter
+
+    Parameters
+    ----------
+    mmc : Core
+
+    Returns
+    -------
+    auto_shutter_state : bool
+    shutter_state : bool
+
+    """
+    auto_shutter_state = mmc.get_auto_shutter()
+    shutter_state = mmc.get_shutter_open()
+
+    return auto_shutter_state, shutter_state
+
+def open_shutter(mmc: Core):
+    """Open shutter if mechanical shutter exists
+
+    Parameters
+    ----------
+    mmc : Core
+
+    """
+
+    if mmc.get_shutter_device():
+        mmc.set_shutter_open(True)
+
+def reset_shutter(mmc: Core, auto_shutter_state: bool, shutter_state: bool):
+    """Reset shutter if mechanical shutter exists
+
+    Parameters
+    ----------
+    mmc : Core
+    auto_shutter_state : bool
+    shutter_state : bool
+
+    """
+
+    if mmc.get_shutter_device():
+        mmc.set_auto_shutter(auto_shutter_state)
+        mmc.set_shutter_open(shutter_state)

From 4a805c476b6d3d9d7c06f500357ea9087bf65394 Mon Sep 17 00:00:00 2001
From: Ivan Ivanov <ivan.ivanov@czbiohub.org>
Date: Fri, 30 Jun 2023 13:26:04 -0700
Subject: [PATCH 10/45] style

---
 mantis/acquisition/microscope_operations.py | 35 ++++++++++++---------
 1 file changed, 20 insertions(+), 15 deletions(-)

diff --git a/mantis/acquisition/microscope_operations.py b/mantis/acquisition/microscope_operations.py
index b9364d96..ad8a3f94 100644
--- a/mantis/acquisition/microscope_operations.py
+++ b/mantis/acquisition/microscope_operations.py
@@ -1,16 +1,15 @@
 import logging
 import time
-import numpy as np
-from pycromanager import Core, Studio
 
-from typing import Tuple, Iterable
+from typing import Iterable, Tuple
 
 import nidaqmx
-
-from nidaqmx.constants import AcquisitionType
+import numpy as np
 
 from copylot.hardware.stages.abstract_stage import AbstractStage
 from copylot.hardware.stages.thorlabs.KIM001 import KCube_PiezoInertia
+from nidaqmx.constants import AcquisitionType
+from pycromanager import Core, Studio
 
 logger = logging.getLogger(__name__)
 
@@ -224,7 +223,8 @@ def autofocus(mmc, mmStudio, z_stage_name: str, z_position):
 
     return autofocus_success
 
-def setup_kim101_stage(serial_number: int, step_rate = 500, step_acceleration = 1000):
+
+def setup_kim101_stage(serial_number: int, step_rate=500, step_acceleration=1000):
     stage = KCube_PiezoInertia(serial_number=str(serial_number))
 
     # Change the acceleration and step rate
@@ -233,6 +233,7 @@ def setup_kim101_stage(serial_number: int, step_rate = 500, step_acceleration =
 
     return stage
 
+
 def create_ram_datastore(
     mmStudio: Studio,
 ):
@@ -241,14 +242,15 @@ def create_ram_datastore(
 
     return datastore
 
+
 def acquire_defocus_stack(
     mmc: Core,
     mmStudio: Studio,
     datastore,
     z_stage,
     z_range: Iterable,
-    channel_ind: int=0,
-    position_ind: int=0,
+    channel_ind: int = 0,
+    position_ind: int = 0,
 ):
     """Snap image at every z position and put image in the datastore
 
@@ -289,8 +291,7 @@ def acquire_defocus_stack(
 
         # get image data
         image_data = np.reshape(
-            tagged_image.pix, 
-            (tagged_image.tags['Height'], tagged_image.tags['Width'])
+            tagged_image.pix, (tagged_image.tags['Height'], tagged_image.tags['Width'])
         )
         data.append(image_data.astype('uint16'))
 
@@ -307,15 +308,16 @@ def acquire_defocus_stack(
 
     return np.asarray(data)
 
+
 def acquire_ls_defocus_stack(
     mmc: Core,
     mmStudio: Studio,
     z_stage,
-    z_start: float, 
-    z_end: float, 
+    z_start: float,
+    z_end: float,
     z_step: float,
-    config_group: str=None, 
-    config_name: str=None,
+    config_group: str = None,
+    config_name: str = None,
 ):
     """Acquire defocus stacks at different galvo positions
 
@@ -344,7 +346,7 @@ def acquire_ls_defocus_stack(
 
     """
     datastore = create_ram_datastore(mmStudio)
-    z_range = np.arange(z_start, z_end+z_step, z_step)
+    z_range = np.arange(z_start, z_end + z_step, z_step)
     data = []
 
     # Set config
@@ -374,6 +376,7 @@ def acquire_ls_defocus_stack(
 
     return np.asarray(data)
 
+
 def get_shutter_state(mmc: Core):
     """Return the current state of the shutter
 
@@ -392,6 +395,7 @@ def get_shutter_state(mmc: Core):
 
     return auto_shutter_state, shutter_state
 
+
 def open_shutter(mmc: Core):
     """Open shutter if mechanical shutter exists
 
@@ -404,6 +408,7 @@ def open_shutter(mmc: Core):
     if mmc.get_shutter_device():
         mmc.set_shutter_open(True)
 
+
 def reset_shutter(mmc: Core, auto_shutter_state: bool, shutter_state: bool):
     """Reset shutter if mechanical shutter exists
 

From cbf624de60c2b7e15847017e0d2859e59dcb72df Mon Sep 17 00:00:00 2001
From: LabelFree <ivan.ivanov@czbiohub.org>
Date: Fri, 30 Jun 2023 17:29:07 -0700
Subject: [PATCH 11/45] update examples

---
 examples/acquire_defocus_stack.py | 24 ++++++++++++++++++------
 examples/test_kim101_stage.py     | 10 ++++------
 2 files changed, 22 insertions(+), 12 deletions(-)

diff --git a/examples/acquire_defocus_stack.py b/examples/acquire_defocus_stack.py
index 807edd96..3936231f 100644
--- a/examples/acquire_defocus_stack.py
+++ b/examples/acquire_defocus_stack.py
@@ -1,13 +1,25 @@
 from pycromanager import Core, Studio
-from mantis.acquisition.microscope_operations import acquire_ls_defocus_stack
+from mantis.acquisition.microscope_operations import setup_kim101_stage, acquire_ls_defocus_stack
 
 mmc = Core()
 mmStudio = Studio()
-z_stage = 'Z'
-z_start = 0
-z_end = 10
-z_step = 1
+z_start = 10
+z_end = 100
+z_step = 10
+config_group = 'Channel - LS'
+config_name = 'GFP EX488 EM525-45'
 
-data = acquire_ls_defocus_stack(mmc, mmStudio, z_stage, z_start, z_end, z_step)
+z_stage = setup_kim101_stage('74000291')
+
+data = acquire_ls_defocus_stack(
+    mmc, 
+    mmStudio, 
+    z_stage, 
+    z_start, 
+    z_end, 
+    z_step, 
+    config_group, 
+    config_name,
+)
 
 print(data.shape)
diff --git a/examples/test_kim101_stage.py b/examples/test_kim101_stage.py
index 24953afa..038778bb 100644
--- a/examples/test_kim101_stage.py
+++ b/examples/test_kim101_stage.py
@@ -15,17 +15,15 @@ def test_light_sheet_stage():
     ### LIGHT SHEET STAGE
     with KCube_PiezoInertia(serial_number='74000291', simulator=False) as stage_LS:
 
-        print(f'LS current position {stage_LS.position}')
-
         # Change the acceleration and step rate
         stage_LS.step_rate = 500
         stage_LS.step_acceleration = 1000
-        print(f'acceleration {stage_LS.step_acceleration} rate {stage_LS.step_rate}')
 
         # Test relative movement
-        step_size = 10
-        stage_LS.move_relative(step_size)
-        stage_LS.move_relative(-step_size)
+        step_size = -10
+        for i in range(10):
+            stage_LS.move_relative(10)
+        # stage_LS.move_relative(-step_size)
 
 
 if __name__ == '__main__':

From 391bf17bbfe6a5cbcfa1edd30b02de954321fda7 Mon Sep 17 00:00:00 2001
From: LabelFree <ivan.ivanov@czbiohub.org>
Date: Fri, 30 Jun 2023 17:29:21 -0700
Subject: [PATCH 12/45] refactor z stage move

---
 mantis/acquisition/microscope_operations.py | 24 ++++++++++++++-------
 1 file changed, 16 insertions(+), 8 deletions(-)

diff --git a/mantis/acquisition/microscope_operations.py b/mantis/acquisition/microscope_operations.py
index b9364d96..bf3e5c25 100644
--- a/mantis/acquisition/microscope_operations.py
+++ b/mantis/acquisition/microscope_operations.py
@@ -4,6 +4,7 @@
 from pycromanager import Core, Studio
 
 from typing import Tuple, Iterable
+from functools import partial
 
 import nidaqmx
 
@@ -272,16 +273,20 @@ def acquire_defocus_stack(
     """
     data = []
 
+    if isinstance(z_stage, str):
+        # this is a MM stage
+        z0 = mmc.get_position(z_stage)
+        move_z = partial(mmc.set_position, z_stage)  # test if this works
+    elif issubclass(type(z_stage), AbstractStage):
+        # this is a copylot stage
+        z0 = z_stage.position
+        move_z = z_stage.move_absolute
+    else:
+        raise RuntimeError(f'Unknown z stage: {z_stage}')
+
     for z_ind, z in enumerate(z_range):
         # set z position
-        if isinstance(z_stage, str):
-            # this is a MM stage
-            mmc.set_relative_position(z_stage, float(z))
-        elif issubclass(z_stage, AbstractStage):
-            # this is a copylot stage
-            z_stage.move_relative(z)
-        else:
-            raise RuntimeError(f'Unknown z stage: {z_stage}')
+        move_z(z0 + z)
 
         # snap image
         mmc.snap_image()
@@ -305,6 +310,9 @@ def acquire_defocus_stack(
         image = image.copy_at_coords(mm_coords)
         datastore.put_image(image)
 
+    # reset z stage
+    move_z(z0)
+    
     return np.asarray(data)
 
 def acquire_ls_defocus_stack(

From 638f384ecf0dc8e0318e9a25894e436fc6006e16 Mon Sep 17 00:00:00 2001
From: LabelFree <ivan.ivanov@czbiohub.org>
Date: Fri, 30 Jun 2023 17:43:49 -0700
Subject: [PATCH 13/45] add galvo scanning to acquire_ls_defocus_stack

---
 mantis/acquisition/microscope_operations.py | 22 ++++++++++++++-------
 1 file changed, 15 insertions(+), 7 deletions(-)

diff --git a/mantis/acquisition/microscope_operations.py b/mantis/acquisition/microscope_operations.py
index 86b5b70d..bbcfe20a 100644
--- a/mantis/acquisition/microscope_operations.py
+++ b/mantis/acquisition/microscope_operations.py
@@ -275,6 +275,7 @@ def acquire_defocus_stack(
     """
     data = []
 
+    # get z0 and define move_z callable for the given stage
     if isinstance(z_stage, str):
         # this is a MM stage
         z0 = mmc.get_position(z_stage)
@@ -321,9 +322,9 @@ def acquire_ls_defocus_stack(
     mmc: Core,
     mmStudio: Studio,
     z_stage,
-    z_start: float,
-    z_end: float,
-    z_step: float,
+    z_range: Iterable,
+    galvo: str,
+    galvo_range: Iterable,
     config_group: str = None,
     config_name: str = None,
 ):
@@ -354,7 +355,6 @@ def acquire_ls_defocus_stack(
 
     """
     datastore = create_ram_datastore(mmStudio)
-    z_range = np.arange(z_start, z_end + z_step, z_step)
     data = []
 
     # Set config
@@ -366,9 +366,13 @@ def acquire_ls_defocus_stack(
     auto_shutter_state, shutter_state = get_shutter_state(mmc)
     open_shutter(mmc)
 
+    # get galvo starting position
+    p0 = mmc.get_position(galvo)
+
     # acquire stack at different galvo positions
-    for p_idx, p in enumerate(range(3)):
-        # TODO: set galvo position
+    for p_idx, p in enumerate(galvo_range):
+        # set galvo position
+        mmc.set_position(galvo, p0 + p)
 
         # acquire defocus stack
         z_stack = acquire_defocus_stack(
@@ -379,6 +383,9 @@ def acquire_ls_defocus_stack(
     # freeze datastore to indicate that we are finished writing to it
     datastore.freeze()
 
+    # Reset galvo
+    mmc.set_position(p0)
+
     # Reset shutter
     reset_shutter(mmc, auto_shutter_state, shutter_state)
 
@@ -414,6 +421,7 @@ def open_shutter(mmc: Core):
     """
 
     if mmc.get_shutter_device():
+        mmc.set_auto_shutter(False)
         mmc.set_shutter_open(True)
 
 
@@ -429,5 +437,5 @@ def reset_shutter(mmc: Core, auto_shutter_state: bool, shutter_state: bool):
     """
 
     if mmc.get_shutter_device():
-        mmc.set_auto_shutter(auto_shutter_state)
         mmc.set_shutter_open(shutter_state)
+        mmc.set_auto_shutter(auto_shutter_state)

From 5721485e1679d669eb0aeda75fea2958575efb33 Mon Sep 17 00:00:00 2001
From: LabelFree <ivan.ivanov@czbiohub.org>
Date: Fri, 30 Jun 2023 17:43:55 -0700
Subject: [PATCH 14/45] Update acquire_defocus_stack.py

---
 examples/acquire_defocus_stack.py | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/examples/acquire_defocus_stack.py b/examples/acquire_defocus_stack.py
index 3936231f..27f0eb21 100644
--- a/examples/acquire_defocus_stack.py
+++ b/examples/acquire_defocus_stack.py
@@ -1,3 +1,4 @@
+import numpy as np
 from pycromanager import Core, Studio
 from mantis.acquisition.microscope_operations import setup_kim101_stage, acquire_ls_defocus_stack
 
@@ -8,16 +9,19 @@
 z_step = 10
 config_group = 'Channel - LS'
 config_name = 'GFP EX488 EM525-45'
+galvo = 'AP Galvo'
+galvo_range = [-1, 0, 1]
 
 z_stage = setup_kim101_stage('74000291')
+z_range = np.arange(z_start, z_end + z_step, z_step)
 
 data = acquire_ls_defocus_stack(
     mmc, 
     mmStudio, 
     z_stage, 
-    z_start, 
-    z_end, 
-    z_step, 
+    z_range,
+    galvo,
+    galvo_range,
     config_group, 
     config_name,
 )

From 4465b496905e650147f8f1c23af3824e30742e07 Mon Sep 17 00:00:00 2001
From: LabelFree <ivan.ivanov@czbiohub.org>
Date: Fri, 7 Jul 2023 09:16:45 -0700
Subject: [PATCH 15/45] update examples

---
 examples/acquire_defocus_stack.py |  6 +++---
 examples/test_kim101_pylablib.py  | 15 +++++++++++++++
 2 files changed, 18 insertions(+), 3 deletions(-)
 create mode 100644 examples/test_kim101_pylablib.py

diff --git a/examples/acquire_defocus_stack.py b/examples/acquire_defocus_stack.py
index 27f0eb21..502dcc88 100644
--- a/examples/acquire_defocus_stack.py
+++ b/examples/acquire_defocus_stack.py
@@ -4,9 +4,9 @@
 
 mmc = Core()
 mmStudio = Studio()
-z_start = 10
-z_end = 100
-z_step = 10
+z_start = -200
+z_end = 200
+z_step = 25
 config_group = 'Channel - LS'
 config_name = 'GFP EX488 EM525-45'
 galvo = 'AP Galvo'
diff --git a/examples/test_kim101_pylablib.py b/examples/test_kim101_pylablib.py
new file mode 100644
index 00000000..547998b4
--- /dev/null
+++ b/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')

From 82ad026d1b56a549134f268b6423db1766bc5107 Mon Sep 17 00:00:00 2001
From: LabelFree <ivan.ivanov@czbiohub.org>
Date: Fri, 7 Jul 2023 10:06:12 -0700
Subject: [PATCH 16/45] switch to pylablib stage and relative moves

---
 mantis/acquisition/microscope_operations.py | 52 ++++++++++++++-------
 pyproject.toml                              |  2 +-
 2 files changed, 37 insertions(+), 17 deletions(-)

diff --git a/mantis/acquisition/microscope_operations.py b/mantis/acquisition/microscope_operations.py
index bbcfe20a..46c41a78 100644
--- a/mantis/acquisition/microscope_operations.py
+++ b/mantis/acquisition/microscope_operations.py
@@ -7,8 +7,7 @@
 import nidaqmx
 import numpy as np
 
-from copylot.hardware.stages.abstract_stage import AbstractStage
-from copylot.hardware.stages.thorlabs.KIM001 import KCube_PiezoInertia
+from pylablib.devices.Thorlabs import KinesisPiezoMotor
 from nidaqmx.constants import AcquisitionType
 from pycromanager import Core, Studio
 
@@ -225,12 +224,29 @@ def autofocus(mmc, mmStudio, z_stage_name: str, z_position):
     return autofocus_success
 
 
-def setup_kim101_stage(serial_number: int, step_rate=500, step_acceleration=1000):
-    stage = KCube_PiezoInertia(serial_number=str(serial_number))
+def setup_kim101_stage(serial_number: int, max_voltage=112, velocity=500, acceleration=1000):
+    """Setup stage on a KIM101 with given drive parameters
 
-    # Change the acceleration and step rate
-    stage.step_rate = step_rate
-    stage.step_acceleration = step_acceleration
+    Parameters
+    ----------
+    serial_number : int
+        8-digit serial number of the KIM101 controller
+    max_voltage : int, optional
+        Max drive voltage in units of Volts, by default 112
+    velocity : int, optional
+        Drive velocity in unit of steps per second, by default 500
+    acceleration : int, optional
+        Drive acceleration in units of steps/s^2, by default 1000
+
+    Returns
+    -------
+    stage : KinesisPiezoMotor
+
+    """
+    stage = KinesisPiezoMotor(str(serial_number))
+
+    # Set drive parameters
+    stage.setup_drive(max_voltage, velocity, acceleration)
 
     return stage
 
@@ -274,22 +290,25 @@ def acquire_defocus_stack(
 
     """
     data = []
+    relative_z_steps = np.hstack((z_range[0], np.diff(z_range)))
+
+    def move_kim101_relative(step):
+        z_stage.move_by(step)
+        z_stage.wait_move()
 
     # get z0 and define move_z callable for the given stage
     if isinstance(z_stage, str):
         # this is a MM stage
-        z0 = mmc.get_position(z_stage)
-        move_z = partial(mmc.set_position, z_stage)  # test if this works
-    elif issubclass(type(z_stage), AbstractStage):
-        # this is a copylot stage
-        z0 = z_stage.position
-        move_z = z_stage.move_absolute
+        move_z = partial(mmc.set_relative_position, z_stage)  # test if this works
+    elif isinstance(z_stage, KinesisPiezoMotor):
+        # this is a pylablib stage
+        move_z = move_kim101_relative
     else:
         raise RuntimeError(f'Unknown z stage: {z_stage}')
 
-    for z_ind, z in enumerate(z_range):
+    for z_ind, rel_z in enumerate(relative_z_steps):
         # set z position
-        move_z(z0 + z)
+        move_z(rel_z)
 
         # snap image
         mmc.snap_image()
@@ -313,7 +332,7 @@ def acquire_defocus_stack(
         datastore.put_image(image)
 
     # reset z stage
-    move_z(z0)
+    move_z(-relative_z_steps.sum())
     
     return np.asarray(data)
 
@@ -329,6 +348,7 @@ def acquire_ls_defocus_stack(
     config_name: str = None,
 ):
     """Acquire defocus stacks at different galvo positions
+    TODO: move to acq_engine.py
 
     Parameters
     ----------
diff --git a/pyproject.toml b/pyproject.toml
index b164937e..5b7802fe 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -30,7 +30,7 @@ dependencies = [
     "pycromanager==0.25.40",
     "nidaqmx",
     "iohub==0.1.0.dev3",
-    "coPylot@git+https://github.com/czbiohub-sf/coPylot.git#egg=thorlabs_stage",
+    "pylablib==1.4.1",
     "matplotlib",
 ]
 

From 94f47b711684860b20243defa47b843514a342f9 Mon Sep 17 00:00:00 2001
From: LabelFree <ivan.ivanov@czbiohub.org>
Date: Fri, 7 Jul 2023 15:44:33 -0700
Subject: [PATCH 17/45] add some logging

---
 examples/acquire_defocus_stack.py           | 9 +++++++++
 mantis/acquisition/microscope_operations.py | 2 ++
 2 files changed, 11 insertions(+)

diff --git a/examples/acquire_defocus_stack.py b/examples/acquire_defocus_stack.py
index 502dcc88..3b74e550 100644
--- a/examples/acquire_defocus_stack.py
+++ b/examples/acquire_defocus_stack.py
@@ -2,6 +2,15 @@
 from pycromanager import Core, Studio
 from mantis.acquisition.microscope_operations import setup_kim101_stage, acquire_ls_defocus_stack
 
+import logging
+logger = logging.getLogger()
+logger.setLevel(logging.DEBUG)
+console_handler = logging.StreamHandler()
+console_handler.setLevel(logging.DEBUG)
+console_format = logging.Formatter('%(levelname)s - %(module)s.%(funcName)s - %(message)s')
+console_handler.setFormatter(console_format)
+logger.addHandler(console_handler)
+
 mmc = Core()
 mmStudio = Studio()
 z_start = -200
diff --git a/mantis/acquisition/microscope_operations.py b/mantis/acquisition/microscope_operations.py
index 46c41a78..d46d9196 100644
--- a/mantis/acquisition/microscope_operations.py
+++ b/mantis/acquisition/microscope_operations.py
@@ -243,9 +243,11 @@ def setup_kim101_stage(serial_number: int, max_voltage=112, velocity=500, accele
     stage : KinesisPiezoMotor
 
     """
+    logger.debug(f'Setting up Kinesis Piezo Motor stage with serial number {serial_number}')
     stage = KinesisPiezoMotor(str(serial_number))
 
     # Set drive parameters
+    logger.debug('Applying drive parameters max voltage: %s, velocity: %s, acceleration: %s', max_voltage, velocity, acceleration)
     stage.setup_drive(max_voltage, velocity, acceleration)
 
     return stage

From f14301785783267b5d52aed3e6b9e85eb356dde7 Mon Sep 17 00:00:00 2001
From: LabelFree <ivan.ivanov@czbiohub.org>
Date: Fri, 7 Jul 2023 17:36:51 -0700
Subject: [PATCH 18/45] add KIM101 compensation

---
 examples/calibrate_kim101.py                | 59 +++++++++++++++++++++
 mantis/acquisition/microscope_operations.py | 28 ++++++++--
 2 files changed, 82 insertions(+), 5 deletions(-)
 create mode 100644 examples/calibrate_kim101.py

diff --git a/examples/calibrate_kim101.py b/examples/calibrate_kim101.py
new file mode 100644
index 00000000..ea20d7aa
--- /dev/null
+++ b/examples/calibrate_kim101.py
@@ -0,0 +1,59 @@
+# 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 
+
+#%%
+import numpy as np
+from pycromanager import Core, Studio
+from mantis.acquisition.microscope_operations import setup_kim101_stage, acquire_ls_defocus_stack
+
+#%%
+mmc = Core()
+mmStudio = Studio()
+z_start = 0
+z_end = 200
+z_step = 25
+galvo = 'AP Galvo'
+galvo_range = [0]*5
+
+z_stage = setup_kim101_stage('74000291')
+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(
+    mmc, 
+    mmStudio, 
+    z_stage, 
+    z_range,
+    galvo,
+    galvo_range,
+)
+
+# %%
+steps_per_direction = len(z_range)//2
+intensity = data.sum(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)
diff --git a/mantis/acquisition/microscope_operations.py b/mantis/acquisition/microscope_operations.py
index d46d9196..ac90633b 100644
--- a/mantis/acquisition/microscope_operations.py
+++ b/mantis/acquisition/microscope_operations.py
@@ -13,6 +13,7 @@
 
 logger = logging.getLogger(__name__)
 
+KIM101_COMPENSATION_FACTOR = 1.08
 
 def _try_mmc_call(mmc, mmc_call_name, *mmc_carr_args):
     """Wrapper that tries to repeat calls to mmCore if they fail. Largely copied
@@ -252,6 +253,27 @@ def setup_kim101_stage(serial_number: int, max_voltage=112, velocity=500, accele
 
     return stage
 
+def set_relative_kim101_position(
+        stage:KinesisPiezoMotor, 
+        step:int, 
+):
+    """Make a relative move with a KIM101 stage, compensating for different 
+    travel distance per step in the positive and negative directions 
+
+    Parameters
+    ----------
+    stage : KinesisPiezoMotor
+        _description_
+    step : int
+        _description_
+    """
+
+    if step < 0:
+        step *= KIM101_COMPENSATION_FACTOR
+
+    stage.move_by(int(step))
+    stage.wait_move()
+
 
 def create_ram_datastore(
     mmStudio: Studio,
@@ -294,17 +316,13 @@ def acquire_defocus_stack(
     data = []
     relative_z_steps = np.hstack((z_range[0], np.diff(z_range)))
 
-    def move_kim101_relative(step):
-        z_stage.move_by(step)
-        z_stage.wait_move()
-
     # get z0 and define move_z callable for the given stage
     if isinstance(z_stage, str):
         # this is a MM stage
         move_z = partial(mmc.set_relative_position, z_stage)  # test if this works
     elif isinstance(z_stage, KinesisPiezoMotor):
         # this is a pylablib stage
-        move_z = move_kim101_relative
+        move_z = partial(set_relative_kim101_position, z_stage)
     else:
         raise RuntimeError(f'Unknown z stage: {z_stage}')
 

From a5ef6577bface22f926a4e022866afb885a58e51 Mon Sep 17 00:00:00 2001
From: LabelFree <ivan.ivanov@czbiohub.org>
Date: Fri, 7 Jul 2023 18:29:18 -0700
Subject: [PATCH 19/45] fix galvo reset position bug

---
 mantis/acquisition/microscope_operations.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mantis/acquisition/microscope_operations.py b/mantis/acquisition/microscope_operations.py
index ac90633b..e6836b2a 100644
--- a/mantis/acquisition/microscope_operations.py
+++ b/mantis/acquisition/microscope_operations.py
@@ -424,7 +424,7 @@ def acquire_ls_defocus_stack(
     datastore.freeze()
 
     # Reset galvo
-    mmc.set_position(p0)
+    mmc.set_position(galvo, p0)
 
     # Reset shutter
     reset_shutter(mmc, auto_shutter_state, shutter_state)

From 276abf65d698f9cdf4276147a003e4365249a576 Mon Sep 17 00:00:00 2001
From: LabelFree <ivan.ivanov@czbiohub.org>
Date: Sun, 9 Jul 2023 15:22:07 -0700
Subject: [PATCH 20/45] add find_focus.py example

---
 examples/acquire_defocus_stack.py |  2 +-
 examples/find_focus.py            | 24 ++++++++++++++++++++++++
 2 files changed, 25 insertions(+), 1 deletion(-)
 create mode 100644 examples/find_focus.py

diff --git a/examples/acquire_defocus_stack.py b/examples/acquire_defocus_stack.py
index 3b74e550..70484e2c 100644
--- a/examples/acquire_defocus_stack.py
+++ b/examples/acquire_defocus_stack.py
@@ -19,7 +19,7 @@
 config_group = 'Channel - LS'
 config_name = 'GFP EX488 EM525-45'
 galvo = 'AP Galvo'
-galvo_range = [-1, 0, 1]
+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)
diff --git a/examples/find_focus.py b/examples/find_focus.py
new file mode 100644
index 00000000..3a00820f
--- /dev/null
+++ b/examples/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()
\ No newline at end of file

From 46e04de13a427d39fce121d334d6c560e424f2fe Mon Sep 17 00:00:00 2001
From: Ivan Ivanov <ivan.ivanov@czbiohub.org>
Date: Sun, 9 Jul 2023 16:23:28 -0700
Subject: [PATCH 21/45] style

---
 mantis/acquisition/microscope_operations.py | 23 ++++++++++++++-------
 1 file changed, 15 insertions(+), 8 deletions(-)

diff --git a/mantis/acquisition/microscope_operations.py b/mantis/acquisition/microscope_operations.py
index e6836b2a..3a2a96cb 100644
--- a/mantis/acquisition/microscope_operations.py
+++ b/mantis/acquisition/microscope_operations.py
@@ -1,20 +1,21 @@
 import logging
 import time
 
-from typing import Iterable, Tuple
 from functools import partial
+from typing import Iterable, Tuple
 
 import nidaqmx
 import numpy as np
 
-from pylablib.devices.Thorlabs import KinesisPiezoMotor
 from nidaqmx.constants import AcquisitionType
 from pycromanager import Core, Studio
+from pylablib.devices.Thorlabs import KinesisPiezoMotor
 
 logger = logging.getLogger(__name__)
 
 KIM101_COMPENSATION_FACTOR = 1.08
 
+
 def _try_mmc_call(mmc, mmc_call_name, *mmc_carr_args):
     """Wrapper that tries to repeat calls to mmCore if they fail. Largely copied
     from dragonfly_automation MicromanagerInterface
@@ -248,17 +249,23 @@ def setup_kim101_stage(serial_number: int, max_voltage=112, velocity=500, accele
     stage = KinesisPiezoMotor(str(serial_number))
 
     # Set drive parameters
-    logger.debug('Applying drive parameters max voltage: %s, velocity: %s, acceleration: %s', max_voltage, velocity, acceleration)
+    logger.debug(
+        'Applying drive parameters max voltage: %s, velocity: %s, acceleration: %s',
+        max_voltage,
+        velocity,
+        acceleration,
+    )
     stage.setup_drive(max_voltage, velocity, acceleration)
 
     return stage
 
+
 def set_relative_kim101_position(
-        stage:KinesisPiezoMotor, 
-        step:int, 
+    stage: KinesisPiezoMotor,
+    step: int,
 ):
-    """Make a relative move with a KIM101 stage, compensating for different 
-    travel distance per step in the positive and negative directions 
+    """Make a relative move with a KIM101 stage, compensating for different
+    travel distance per step in the positive and negative directions
 
     Parameters
     ----------
@@ -353,7 +360,7 @@ def acquire_defocus_stack(
 
     # reset z stage
     move_z(-relative_z_steps.sum())
-    
+
     return np.asarray(data)
 
 

From 6372a18199011657a8d0c141421c3fac2a846477 Mon Sep 17 00:00:00 2001
From: Ivan Ivanov <ivan.ivanov@czbiohub.org>
Date: Mon, 10 Jul 2023 13:26:25 -0700
Subject: [PATCH 22/45] ignore tests in examples dir

---
 pyproject.toml | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/pyproject.toml b/pyproject.toml
index 5b7802fe..804cf18b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -109,3 +109,6 @@ skip_glob = ["examples/*"]
 disable = ["C", "R", "W", "unsubscriptable-object", "import-error"]
 msg-template = "{line},{column},{category},{symbol}:{msg}"
 reports = "n"
+
+[tool.pytest.ini_options]
+addopts = "--ignore examples/"

From 702e8f368e100cda9a32327d42f02ae736a0b7a3 Mon Sep 17 00:00:00 2001
From: LabelFree <ivan.ivanov@czbiohub.org>
Date: Mon, 10 Jul 2023 18:04:30 -0700
Subject: [PATCH 23/45] move o3 refocus to acq_engine

---
 examples/acquire_defocus_stack.py           |   3 +-
 mantis/acquisition/AcquisitionSettings.py   |   2 +
 mantis/acquisition/acq_engine.py            | 131 ++++++++++++++++++++
 mantis/acquisition/microscope_operations.py |  96 +++-----------
 4 files changed, 154 insertions(+), 78 deletions(-)

diff --git a/examples/acquire_defocus_stack.py b/examples/acquire_defocus_stack.py
index 70484e2c..31e9dcfa 100644
--- a/examples/acquire_defocus_stack.py
+++ b/examples/acquire_defocus_stack.py
@@ -1,6 +1,7 @@
 import numpy as np
 from pycromanager import Core, Studio
-from mantis.acquisition.microscope_operations import setup_kim101_stage, acquire_ls_defocus_stack
+from mantis.acquisition.microscope_operations import setup_kim101_stage
+from mantis.acquisition.acq_engine import acquire_ls_defocus_stack
 
 import logging
 logger = logging.getLogger()
diff --git a/mantis/acquisition/AcquisitionSettings.py b/mantis/acquisition/AcquisitionSettings.py
index cfc7b16c..b16c5924 100644
--- a/mantis/acquisition/AcquisitionSettings.py
+++ b/mantis/acquisition/AcquisitionSettings.py
@@ -80,3 +80,5 @@ 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[Tuple[str, str]] = None
diff --git a/mantis/acquisition/acq_engine.py b/mantis/acquisition/acq_engine.py
index 0e47a9b3..829cb26b 100644
--- a/mantis/acquisition/acq_engine.py
+++ b/mantis/acquisition/acq_engine.py
@@ -1,6 +1,7 @@
 import logging
 import os
 import time
+from typing import Iterable
 
 from copy import deepcopy
 from dataclasses import asdict
@@ -16,6 +17,8 @@
 from mantis.acquisition import microscope_operations
 from mantis.acquisition.logger import configure_logger
 
+from waveorder.focus import focus_from_transverse_band
+
 # isort: off
 from mantis.acquisition.AcquisitionSettings import (
     TimeSettings,
@@ -50,6 +53,9 @@
 LS_KIM101_SN = 74000291
 LF_KIM101_SN = 74000565
 
+NA_DETECTION = 1.35
+PIXEL_SIZE = 6.5 / (40 * 1.4)  # in um
+
 logger = logging.getLogger(__name__)
 
 
@@ -87,6 +93,7 @@ def __init__(
         self.type = 'light-sheet' if self.headless else 'label-free'
         self.mmc = None
         self.mmStudio = None
+        self.o3_stage = None
 
         logger.debug(f'Initializing {self.type} acquisition engine')
         if enabled:
@@ -185,6 +192,10 @@ def setup(self):
                 self.mmc, 'Core', 'Focus', self.slice_settings.z_stage_name
             )
 
+            # Setup O3 scan stage
+            if self.microscope_settings.use_o3_refocus:
+                self.o3_stage = microscope_operations.setup_kim101_stage(LS_KIM101_SN)
+
             # Note: sequencing should be turned off by default
             # Setup z sequencing
             if self.slice_settings.use_sequencing:
@@ -592,6 +603,52 @@ def go_to_position(self, position_index: int):
                     self.lf_acq.microscope_settings.autofocus_stage,
                 )
 
+    def refocus_ls_path(self):
+        # Define O3 z range
+        # 1 step is approx 20 nm, 25 steps are 500 nm which is approx Nyquist sampling
+        z_start = -200
+        z_end = 200
+        z_step = 25
+        z_range = np.arange(z_start, z_end + z_step, z_step)
+
+        # Define galvo range, i.e. galvo positions at which O3 defocus stacks
+        # are acquired, should be odd number
+        galvo_scan_range = self.ls_acq.slice_settings.z_start
+        len_galvo_scan_range = len(galvo_scan_range)
+        galvo_range = [
+            galvo_scan_range[int(0.3*len_galvo_scan_range)],
+            galvo_scan_range[int(0.5*len_galvo_scan_range)],
+            galvo_scan_range[int(0.7*len_galvo_scan_range)],
+        ]
+        
+        # Acquire defocus stacks at several galvo positions
+        data = acquire_ls_defocus_stack(
+            mmc=self.ls_acq.mmc,
+            mmStudio=self.ls_acq.mmStudio,
+            z_stage=self.ls_acq.o3_stage,
+            z_range=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[0],
+            config_name=self.ls_acq.microscope_settings.o3_refocus_config[1],
+        )
+
+        # Find in-focus slice
+        wavelength = 0.55  # in um, approx
+        focus_indices = []
+        for stack in data:
+            idx = focus_from_transverse_band(
+                stack, NA_det=NA_DETECTION, lambda_ill=wavelength, pixel_size=PIXEL_SIZE
+            )
+            focus_indices.append(idx)
+        
+        # Refocus O3
+        focus_idx = int(np.median(focus_indices))
+        microscope_operations.set_relative_kim101_position(
+            self.ls_acq.o3_stage,
+            z_range[focus_idx]
+        )
+    
     def setup(self):
         """
         Setup the mantis acquisition. This method sets up the label-free
@@ -812,3 +869,77 @@ def _create_acquisition_directory(root_dir, acq_name, idx=1):
     except OSError:
         return _create_acquisition_directory(root_dir, acq_name, idx + 1)
     return acq_dir
+
+
+def acquire_ls_defocus_stack(
+    mmc: Core,
+    mmStudio: Studio,
+    z_stage,
+    z_range: Iterable,
+    galvo: str,
+    galvo_range: Iterable,
+    config_group: str = None,
+    config_name: str = None,
+):
+    """Acquire defocus stacks at different galvo positions
+
+    Parameters
+    ----------
+    mmc : Core
+        _description_
+    mmStudio : Studio
+        _description_
+    z_stage : _type_
+        _description_
+    z_start : float
+        _description_
+    z_end : float
+        _description_
+    z_step : float
+        _description_
+    config_group : str, optional
+        _description_, by default None
+    config_name : str, optional
+        _description_, by default None
+
+    Returns
+    -------
+    data : np.array
+
+    """
+    datastore = microscope_operations.create_ram_datastore(mmStudio)
+    data = []
+
+    # Set config
+    if config_name is not None:
+        mmc.set_config(config_group, config_name)
+        mmc.wait_for_config(config_group, config_name)
+
+    # Open shutter
+    auto_shutter_state, shutter_state = microscope_operations.get_shutter_state(mmc)
+    microscope_operations.open_shutter(mmc)
+
+    # get galvo starting position
+    p0 = mmc.get_position(galvo)
+
+    # acquire stack at different galvo positions
+    for p_idx, p in enumerate(galvo_range):
+        # set galvo position
+        mmc.set_position(galvo, p0 + p)
+
+        # acquire defocus stack
+        z_stack = microscope_operations.acquire_defocus_stack(
+            mmc, mmStudio, datastore, z_stage, z_range, channel_ind=0, position_ind=p_idx
+        )
+        data.append(z_stack)
+
+    # freeze datastore to indicate that we are finished writing to it
+    datastore.freeze()
+
+    # Reset galvo
+    mmc.set_position(galvo, p0)
+
+    # Reset shutter
+    microscope_operations.reset_shutter(mmc, auto_shutter_state, shutter_state)
+
+    return np.asarray(data)
diff --git a/mantis/acquisition/microscope_operations.py b/mantis/acquisition/microscope_operations.py
index 3a2a96cb..0973245f 100644
--- a/mantis/acquisition/microscope_operations.py
+++ b/mantis/acquisition/microscope_operations.py
@@ -257,6 +257,10 @@ def setup_kim101_stage(serial_number: int, max_voltage=112, velocity=500, accele
     )
     stage.setup_drive(max_voltage, velocity, acceleration)
 
+    # manually keep track of the stage true position based on the relative
+    # number of steps in each direction
+    stage.true_position = 0
+
     return stage
 
 
@@ -275,6 +279,10 @@ def set_relative_kim101_position(
         _description_
     """
 
+    # keep track of the stage actual position, in steps, not accounting for the
+    # compensation factor
+    stage.true_position += step
+
     if step < 0:
         step *= KIM101_COMPENSATION_FACTOR
 
@@ -364,81 +372,6 @@ def acquire_defocus_stack(
     return np.asarray(data)
 
 
-def acquire_ls_defocus_stack(
-    mmc: Core,
-    mmStudio: Studio,
-    z_stage,
-    z_range: Iterable,
-    galvo: str,
-    galvo_range: Iterable,
-    config_group: str = None,
-    config_name: str = None,
-):
-    """Acquire defocus stacks at different galvo positions
-    TODO: move to acq_engine.py
-
-    Parameters
-    ----------
-    mmc : Core
-        _description_
-    mmStudio : Studio
-        _description_
-    z_stage : _type_
-        _description_
-    z_start : float
-        _description_
-    z_end : float
-        _description_
-    z_step : float
-        _description_
-    config_group : str, optional
-        _description_, by default None
-    config_name : str, optional
-        _description_, by default None
-
-    Returns
-    -------
-    data : np.array
-
-    """
-    datastore = create_ram_datastore(mmStudio)
-    data = []
-
-    # Set config
-    if config_name is not None:
-        mmc.set_config(config_group, config_name)
-        mmc.wait_for_config(config_group, config_name)
-
-    # Open shutter
-    auto_shutter_state, shutter_state = get_shutter_state(mmc)
-    open_shutter(mmc)
-
-    # get galvo starting position
-    p0 = mmc.get_position(galvo)
-
-    # acquire stack at different galvo positions
-    for p_idx, p in enumerate(galvo_range):
-        # set galvo position
-        mmc.set_position(galvo, p0 + p)
-
-        # acquire defocus stack
-        z_stack = acquire_defocus_stack(
-            mmc, mmStudio, datastore, z_stage, z_range, channel_ind=0, position_ind=p_idx
-        )
-        data.append(z_stack)
-
-    # freeze datastore to indicate that we are finished writing to it
-    datastore.freeze()
-
-    # Reset galvo
-    mmc.set_position(galvo, p0)
-
-    # Reset shutter
-    reset_shutter(mmc, auto_shutter_state, shutter_state)
-
-    return np.asarray(data)
-
-
 def get_shutter_state(mmc: Core):
     """Return the current state of the shutter
 
@@ -467,7 +400,9 @@ def open_shutter(mmc: Core):
 
     """
 
-    if mmc.get_shutter_device():
+    shutter_device = mmc.get_shutter_device()
+    if shutter_device:
+        logger.debug(f'Opening shutter {shutter_device}')
         mmc.set_auto_shutter(False)
         mmc.set_shutter_open(True)
 
@@ -483,6 +418,13 @@ def reset_shutter(mmc: Core, auto_shutter_state: bool, shutter_state: bool):
 
     """
 
-    if mmc.get_shutter_device():
+    shutter_device = mmc.get_shutter_device()
+    if shutter_device:
+        logger.debug(
+            'Resetting shutter %s to state Open:%s, Autoshutter: %s', 
+            shutter_device, 
+            shutter_state, 
+            auto_shutter_state
+        )
         mmc.set_shutter_open(shutter_state)
         mmc.set_auto_shutter(auto_shutter_state)

From 211b49d11e5fa327c5a4e6c19f6080fa9767e146 Mon Sep 17 00:00:00 2001
From: LabelFree <ivan.ivanov@czbiohub.org>
Date: Tue, 11 Jul 2023 17:17:58 -0700
Subject: [PATCH 24/45] add checks and logging

---
 examples/acquire_defocus_stack.py |  1 +
 mantis/acquisition/acq_engine.py  | 34 ++++++++++++++++++++++++-------
 2 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/examples/acquire_defocus_stack.py b/examples/acquire_defocus_stack.py
index 31e9dcfa..45ecd88a 100644
--- a/examples/acquire_defocus_stack.py
+++ b/examples/acquire_defocus_stack.py
@@ -34,6 +34,7 @@
     galvo_range,
     config_group, 
     config_name,
+    close_display=False,
 )
 
 print(data.shape)
diff --git a/mantis/acquisition/acq_engine.py b/mantis/acquisition/acq_engine.py
index 829cb26b..dad057e6 100644
--- a/mantis/acquisition/acq_engine.py
+++ b/mantis/acquisition/acq_engine.py
@@ -54,7 +54,7 @@
 LF_KIM101_SN = 74000565
 
 NA_DETECTION = 1.35
-PIXEL_SIZE = 6.5 / (40 * 1.4)  # in um
+LS_PIXEL_SIZE = 6.5 / (40 * 1.4)  # in um
 
 logger = logging.getLogger(__name__)
 
@@ -604,6 +604,8 @@ def go_to_position(self, position_index: int):
                 )
 
     def refocus_ls_path(self):
+        logger.info('Running O3 refocus algorithm on light-sheet arm')
+
         # Define O3 z range
         # 1 step is approx 20 nm, 25 steps are 500 nm which is approx Nyquist sampling
         z_start = -200
@@ -631,6 +633,7 @@ def refocus_ls_path(self):
             galvo_range=galvo_range,
             config_group=self.ls_acq.microscope_settings.o3_refocus_config[0],
             config_name=self.ls_acq.microscope_settings.o3_refocus_config[1],
+            close_display=True,
         )
 
         # Find in-focus slice
@@ -638,16 +641,25 @@ def refocus_ls_path(self):
         focus_indices = []
         for stack in data:
             idx = focus_from_transverse_band(
-                stack, NA_det=NA_DETECTION, lambda_ill=wavelength, pixel_size=PIXEL_SIZE
+                stack, NA_det=NA_DETECTION, lambda_ill=wavelength, pixel_size=LS_PIXEL_SIZE
             )
             focus_indices.append(idx)
+        logger.debug(f'Stacks at galvo positions {galvo_range} are in focus at slice {focus_indices}')
         
         # Refocus O3
-        focus_idx = int(np.median(focus_indices))
-        microscope_operations.set_relative_kim101_position(
-            self.ls_acq.o3_stage,
-            z_range[focus_idx]
-        )
+        # Some focus_indices may be None, e.g. if there is no sample
+        valid_focus_indices = [idx for idx in focus_indices if idx is not None]
+        if valid_focus_indices:
+            focus_idx = int(np.median(valid_focus_indices))
+            o3_displacement = int(z_range[focus_idx])
+
+            logger.info(f'Moving O3 by {o3_displacement} steps')
+            microscope_operations.set_relative_kim101_position(
+                self.ls_acq.o3_stage,
+                o3_displacement
+            )
+        else:
+            logger.error('Could not determine the correct O3 in-focus position. O3 will not move')
     
     def setup(self):
         """
@@ -880,6 +892,7 @@ def acquire_ls_defocus_stack(
     galvo_range: Iterable,
     config_group: str = None,
     config_name: str = None,
+    close_display: bool = True,
 ):
     """Acquire defocus stacks at different galvo positions
 
@@ -901,6 +914,8 @@ def acquire_ls_defocus_stack(
         _description_, by default None
     config_name : str, optional
         _description_, by default None
+    close_display: bool. optional
+        _description_, by default True
 
     Returns
     -------
@@ -942,4 +957,9 @@ def acquire_ls_defocus_stack(
     # Reset shutter
     microscope_operations.reset_shutter(mmc, auto_shutter_state, shutter_state)
 
+    # Close datastore and associated displays; if close_display=False, display
+    # window must be manually closed
+    if close_display:
+        datastore.close()
+
     return np.asarray(data)

From b3b34e74afdc55f7499d3739a97ed86cf18cb552 Mon Sep 17 00:00:00 2001
From: LabelFree <ivan.ivanov@czbiohub.org>
Date: Wed, 12 Jul 2023 12:05:51 -0700
Subject: [PATCH 25/45] debug acq engine o3 refocus

---
 mantis/acquisition/AcquisitionSettings.py   |   3 +-
 mantis/acquisition/acq_engine.py            |  40 +++----
 mantis/acquisition/microscope_operations.py | 110 +++++++++++++++++---
 3 files changed, 120 insertions(+), 33 deletions(-)

diff --git a/mantis/acquisition/AcquisitionSettings.py b/mantis/acquisition/AcquisitionSettings.py
index b16c5924..d1431bc0 100644
--- a/mantis/acquisition/AcquisitionSettings.py
+++ b/mantis/acquisition/AcquisitionSettings.py
@@ -81,4 +81,5 @@ class MicroscopeSettings:
     autofocus_stage: Optional[str] = None
     autofocus_method: Optional[str] = None
     use_o3_refocus: bool = False
-    o3_refocus_config: Optional[Tuple[str, str]] = None
+    o3_refocus_config: Optional[ConfigSettings] = None
+    o3_refocus_interval_min: Optional[int] = None
diff --git a/mantis/acquisition/acq_engine.py b/mantis/acquisition/acq_engine.py
index dad057e6..81907817 100644
--- a/mantis/acquisition/acq_engine.py
+++ b/mantis/acquisition/acq_engine.py
@@ -608,6 +608,7 @@ def refocus_ls_path(self):
 
         # Define O3 z range
         # 1 step is approx 20 nm, 25 steps are 500 nm which is approx Nyquist sampling
+        # The stack starts away from O2 and moves closer
         z_start = -200
         z_end = 200
         z_step = 25
@@ -615,7 +616,7 @@ def refocus_ls_path(self):
 
         # Define galvo range, i.e. galvo positions at which O3 defocus stacks
         # are acquired, should be odd number
-        galvo_scan_range = self.ls_acq.slice_settings.z_start
+        galvo_scan_range = self.ls_acq.slice_settings.z_range
         len_galvo_scan_range = len(galvo_scan_range)
         galvo_range = [
             galvo_scan_range[int(0.3*len_galvo_scan_range)],
@@ -626,14 +627,12 @@ def refocus_ls_path(self):
         # Acquire defocus stacks at several galvo positions
         data = acquire_ls_defocus_stack(
             mmc=self.ls_acq.mmc,
-            mmStudio=self.ls_acq.mmStudio,
             z_stage=self.ls_acq.o3_stage,
             z_range=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[0],
-            config_name=self.ls_acq.microscope_settings.o3_refocus_config[1],
-            close_display=True,
+            config_group=self.ls_acq.microscope_settings.o3_refocus_config.config_group,
+            config_name=self.ls_acq.microscope_settings.o3_refocus_config.config_name,
         )
 
         # Find in-focus slice
@@ -644,7 +643,11 @@ def refocus_ls_path(self):
                 stack, NA_det=NA_DETECTION, lambda_ill=wavelength, pixel_size=LS_PIXEL_SIZE
             )
             focus_indices.append(idx)
-        logger.debug(f'Stacks at galvo positions {galvo_range} are in focus at slice {focus_indices}')
+        logger.debug(
+            'Stacks at galvo positions %s are in focus at slice %s', 
+            np.round(galvo_range, 3), 
+            focus_indices
+        )
         
         # Refocus O3
         # Some focus_indices may be None, e.g. if there is no sample
@@ -788,6 +791,11 @@ 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:
+                    self.refocus_ls_path()
+
                 # start acquisition
                 lf_events = deepcopy(lf_cz_events)
                 for _event in lf_events:
@@ -885,14 +893,12 @@ def _create_acquisition_directory(root_dir, acq_name, idx=1):
 
 def acquire_ls_defocus_stack(
     mmc: Core,
-    mmStudio: Studio,
     z_stage,
     z_range: Iterable,
     galvo: str,
     galvo_range: Iterable,
     config_group: str = None,
     config_name: str = None,
-    close_display: bool = True,
 ):
     """Acquire defocus stacks at different galvo positions
 
@@ -922,7 +928,6 @@ def acquire_ls_defocus_stack(
     data : np.array
 
     """
-    datastore = microscope_operations.create_ram_datastore(mmStudio)
     data = []
 
     # Set config
@@ -937,19 +942,21 @@ def acquire_ls_defocus_stack(
     # get galvo starting position
     p0 = mmc.get_position(galvo)
 
+    # set camera to internal trigger
+    # TODO: do this properly, context manager?
+    microscope_operations.set_property(mmc, 'Prime BSI Express', 'TriggerMode', 'Internal Trigger')
+
     # acquire stack at different galvo positions
     for p_idx, p in enumerate(galvo_range):
         # set galvo position
         mmc.set_position(galvo, p0 + p)
 
         # acquire defocus stack
-        z_stack = microscope_operations.acquire_defocus_stack(
-            mmc, mmStudio, datastore, z_stage, z_range, channel_ind=0, position_ind=p_idx
-        )
+        z_stack = microscope_operations.acquire_defocus_stack(mmc, z_stage, z_range)
         data.append(z_stack)
 
-    # freeze datastore to indicate that we are finished writing to it
-    datastore.freeze()
+    # Reset camera triggering
+    microscope_operations.set_property(mmc, 'Prime BSI Express', 'TriggerMode', 'Edge Trigger')
 
     # Reset galvo
     mmc.set_position(galvo, p0)
@@ -957,9 +964,4 @@ def acquire_ls_defocus_stack(
     # Reset shutter
     microscope_operations.reset_shutter(mmc, auto_shutter_state, shutter_state)
 
-    # Close datastore and associated displays; if close_display=False, display
-    # window must be manually closed
-    if close_display:
-        datastore.close()
-
     return np.asarray(data)
diff --git a/mantis/acquisition/microscope_operations.py b/mantis/acquisition/microscope_operations.py
index 0973245f..449e7b3a 100644
--- a/mantis/acquisition/microscope_operations.py
+++ b/mantis/acquisition/microscope_operations.py
@@ -301,10 +301,10 @@ def create_ram_datastore(
 
 def acquire_defocus_stack(
     mmc: Core,
-    mmStudio: Studio,
-    datastore,
     z_stage,
     z_range: Iterable,
+    mmStudio: Studio = None,
+    datastore = None,
     channel_ind: int = 0,
     position_ind: int = 0,
 ):
@@ -313,11 +313,12 @@ def acquire_defocus_stack(
     Parameters
     ----------
     mmc : Core
-    mmStudio : Studio
-    datastore : micromanager.data.Datastore
-        Micro-manager datastore object
     z_stage : str or coPylot stage object
     z_range : Iterable
+    mmStudio : Studio, optional
+        If not None, images will be added to a Micro-manager RAM datastore
+    datastore : micromanager.data.Datastore
+        Micro-manager datastore object
     channel_ind : int, optional
         Channel index of acquired images in the Micro-manager datastore, by default 0
     position_ind : int, optional
@@ -341,6 +342,7 @@ def acquire_defocus_stack(
     else:
         raise RuntimeError(f'Unknown z stage: {z_stage}')
 
+    # mmc.wait_for_image_synchro()
     for z_ind, rel_z in enumerate(relative_z_steps):
         # set z position
         move_z(rel_z)
@@ -356,21 +358,103 @@ def acquire_defocus_stack(
         data.append(image_data.astype('uint16'))
 
         # set image coordinates and put in datastore
-        image = mmStudio.get_data_manager().convert_tagged_image(tagged_image)
-        coords_builder = image.get_coords().copy()
-        coords_builder = coords_builder.z(z_ind)
-        coords_builder = coords_builder.channel(channel_ind)
-        coords_builder = coords_builder.stage_position(position_ind)
-        mm_coords = coords_builder.build()
+        if mmStudio is not None:
+            image = mmStudio.get_data_manager().convert_tagged_image(tagged_image)
+            coords_builder = image.get_coords().copy()
+            coords_builder = coords_builder.z(z_ind)
+            coords_builder = coords_builder.channel(channel_ind)
+            coords_builder = coords_builder.stage_position(position_ind)
+            mm_coords = coords_builder.build()
 
-        image = image.copy_at_coords(mm_coords)
-        datastore.put_image(image)
+            image = image.copy_at_coords(mm_coords)
+            datastore.put_image(image)
 
     # reset z stage
     move_z(-relative_z_steps.sum())
 
     return np.asarray(data)
 
+def acquire_ls_defocus_stack_and_display(
+    mmc: Core,
+    mmStudio: Studio,
+    z_stage,
+    z_range: Iterable,
+    galvo: str,
+    galvo_range: Iterable,
+    config_group: str = None,
+    config_name: str = None,
+    close_display: bool = True,
+):
+    """Acquire defocus stacks at different galvo positions
+
+    Parameters
+    ----------
+    mmc : Core
+        _description_
+    mmStudio : Studio
+        _description_
+    z_stage : _type_
+        _description_
+    z_start : float
+        _description_
+    z_end : float
+        _description_
+    z_step : float
+        _description_
+    config_group : str, optional
+        _description_, by default None
+    config_name : str, optional
+        _description_, by default None
+    close_display: bool. optional
+        _description_, by default True
+
+    Returns
+    -------
+    data : np.array
+
+    """
+    datastore = create_ram_datastore(mmStudio)
+    data = []
+
+    # Set config
+    if config_name is not None:
+        mmc.set_config(config_group, config_name)
+        mmc.wait_for_config(config_group, config_name)
+
+    # Open shutter
+    auto_shutter_state, shutter_state = get_shutter_state(mmc)
+    open_shutter(mmc)
+
+    # get galvo starting position
+    p0 = mmc.get_position(galvo)
+
+    # acquire stack at different galvo positions
+    for p_idx, p in enumerate(galvo_range):
+        # set galvo position
+        mmc.set_position(galvo, p0 + p)
+
+        # acquire defocus stack
+        z_stack = acquire_defocus_stack(
+            mmc, z_stage, z_range, mmStudio, datastore, channel_ind=0, position_ind=p_idx
+        )
+        data.append(z_stack)
+
+    # freeze datastore to indicate that we are finished writing to it
+    datastore.freeze()
+
+    # Reset galvo
+    mmc.set_position(galvo, p0)
+
+    # Reset shutter
+    reset_shutter(mmc, auto_shutter_state, shutter_state)
+
+    # Close datastore and associated displays; if close_display=False, display
+    # window must be manually closed
+    if close_display:
+        datastore.close()
+
+    return np.asarray(data)
+
 
 def get_shutter_state(mmc: Core):
     """Return the current state of the shutter

From 838e77b4fa3cdc242aa67a5e25d103ed05141075 Mon Sep 17 00:00:00 2001
From: LabelFree <ivan.ivanov@czbiohub.org>
Date: Wed, 12 Jul 2023 12:06:09 -0700
Subject: [PATCH 26/45] update example

---
 examples/acquire_defocus_stack.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/examples/acquire_defocus_stack.py b/examples/acquire_defocus_stack.py
index 45ecd88a..a58a70ad 100644
--- a/examples/acquire_defocus_stack.py
+++ b/examples/acquire_defocus_stack.py
@@ -1,7 +1,6 @@
 import numpy as np
 from pycromanager import Core, Studio
-from mantis.acquisition.microscope_operations import setup_kim101_stage
-from mantis.acquisition.acq_engine import acquire_ls_defocus_stack
+from mantis.acquisition.microscope_operations import setup_kim101_stage, acquire_ls_defocus_stack_and_display
 
 import logging
 logger = logging.getLogger()
@@ -25,7 +24,7 @@
 z_stage = setup_kim101_stage('74000291')
 z_range = np.arange(z_start, z_end + z_step, z_step)
 
-data = acquire_ls_defocus_stack(
+data = acquire_ls_defocus_stack_and_display(
     mmc, 
     mmStudio, 
     z_stage, 

From 9177c872e1c8a965c25cccf5eabbdf59a4abbd38 Mon Sep 17 00:00:00 2001
From: LabelFree <ivan.ivanov@czbiohub.org>
Date: Wed, 12 Jul 2023 12:11:28 -0700
Subject: [PATCH 27/45] Create separate logs directory

---
 mantis/acquisition/acq_engine.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/mantis/acquisition/acq_engine.py b/mantis/acquisition/acq_engine.py
index 81907817..9352b257 100644
--- a/mantis/acquisition/acq_engine.py
+++ b/mantis/acquisition/acq_engine.py
@@ -294,13 +294,15 @@ def __init__(
         if not enable_lf_acq or not enable_ls_acq:
             raise Exception('Disabling LF or LS acquisition is not currently supported')
 
-        # Create acquisition directory
+        # Create acquisition directory and log directory
         self._acq_dir = _create_acquisition_directory(self._root_dir, self._acq_name)
+        self._logs_dir = os.path.join(self._acq_dir, 'logs')
+        os.mkdir(self._logs_dir)
 
         # Setup logger
         timestamp = datetime.now().strftime("%Y%m%dT%H%M%S")
         configure_logger(
-            os.path.join(self._acq_dir, f'mantis_acquisition_log_{timestamp}.txt')
+            os.path.join(self._logs_dir, f'acquisition_log_{timestamp}.txt')
         )
 
         # initialize time and position settings

From 6dd197f3db86c3f54856e09eed47395d7ac62781 Mon Sep 17 00:00:00 2001
From: LabelFree <ivan.ivanov@czbiohub.org>
Date: Wed, 12 Jul 2023 12:15:59 -0700
Subject: [PATCH 28/45] move conda env logger to logs

---
 mantis/acquisition/acq_engine.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mantis/acquisition/acq_engine.py b/mantis/acquisition/acq_engine.py
index caf7e5d9..16a35f03 100644
--- a/mantis/acquisition/acq_engine.py
+++ b/mantis/acquisition/acq_engine.py
@@ -315,7 +315,7 @@ def __init__(
 
         # Log conda environment
         outs, errs = log_conda_environment(
-            os.path.join(self._acq_dir, f'conda_environment_log_{timestamp}.txt')
+            os.path.join(self._logs_dir, f'conda_environment_log_{timestamp}.txt')
         )
         if errs is None:
             logger.debug(outs.decode('ascii').strip())

From 1b7ed37c4e4c4b268775fc4cf361124ec0957a75 Mon Sep 17 00:00:00 2001
From: LabelFree <ivan.ivanov@czbiohub.org>
Date: Wed, 12 Jul 2023 12:20:15 -0700
Subject: [PATCH 29/45] rename to mantis_acquisition_log

---
 mantis/acquisition/acq_engine.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mantis/acquisition/acq_engine.py b/mantis/acquisition/acq_engine.py
index 16a35f03..ff215b0a 100644
--- a/mantis/acquisition/acq_engine.py
+++ b/mantis/acquisition/acq_engine.py
@@ -302,7 +302,7 @@ def __init__(
         # Setup logger
         timestamp = datetime.now().strftime("%Y%m%dT%H%M%S")
         configure_logger(
-            os.path.join(self._logs_dir, f'acquisition_log_{timestamp}.txt')
+            os.path.join(self._logs_dir, f'mantis_acquisition_log_{timestamp}.txt')
         )
 
         # initialize time and position settings

From 136bd724731d4c4d544ed16ca72e4ba45a2d80da Mon Sep 17 00:00:00 2001
From: LabelFree <ivan.ivanov@czbiohub.org>
Date: Wed, 12 Jul 2023 16:57:10 -0700
Subject: [PATCH 30/45] save acquired stacks

---
 mantis/acquisition/acq_engine.py | 8 ++++++++
 pyproject.toml                   | 1 +
 2 files changed, 9 insertions(+)

diff --git a/mantis/acquisition/acq_engine.py b/mantis/acquisition/acq_engine.py
index ff215b0a..fec7dc02 100644
--- a/mantis/acquisition/acq_engine.py
+++ b/mantis/acquisition/acq_engine.py
@@ -8,6 +8,7 @@
 from datetime import datetime
 from functools import partial
 
+import tifffile
 import nidaqmx
 import numpy as np
 
@@ -646,6 +647,13 @@ def refocus_ls_path(self):
             config_name=self.ls_acq.microscope_settings.o3_refocus_config.config_name,
         )
 
+        # Save acquired stacks
+        timestamp = datetime.now().strftime("%Y%m%dT%H%M%S")
+        tifffile.imwrite(
+            os.path.join(self._logs_dir, f'ls_refocus_data_{timestamp}.ome.tif'),
+            np.expand_dims(data, -3).astype('uint16')
+        )
+
         # Find in-focus slice
         wavelength = 0.55  # in um, approx
         focus_indices = []
diff --git a/pyproject.toml b/pyproject.toml
index 319816d9..632e3eab 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -33,6 +33,7 @@ dependencies = [
     "scipy",
     "slurmkit @ git+https://github.com/royerlab/slurmkit",
     "natsort",
+    "tifffile",
     "napari[all]; 'arm64' not in platform_machine", # with Qt5 and skimage
     "napari; 'arm64' in platform_machine", # without Qt5 and skimage
     "PyQt6; 'arm64' in platform_machine",

From 106609e19b0bd36fbd4d905600f2b8642c38d218 Mon Sep 17 00:00:00 2001
From: LabelFree <ivan.ivanov@czbiohub.org>
Date: Wed, 12 Jul 2023 17:17:44 -0700
Subject: [PATCH 31/45] implement timed o3 refocus

---
 mantis/acquisition/acq_engine.py | 20 +++++++++++++-------
 1 file changed, 13 insertions(+), 7 deletions(-)

diff --git a/mantis/acquisition/acq_engine.py b/mantis/acquisition/acq_engine.py
index fec7dc02..1af3b4a5 100644
--- a/mantis/acquisition/acq_engine.py
+++ b/mantis/acquisition/acq_engine.py
@@ -619,11 +619,11 @@ def refocus_ls_path(self):
         logger.info('Running O3 refocus algorithm on light-sheet arm')
 
         # Define O3 z range
-        # 1 step is approx 20 nm, 25 steps are 500 nm which is approx Nyquist sampling
+        # 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
-        z_start = -200
-        z_end = 200
-        z_step = 25
+        z_start = -105
+        z_end = 105
+        z_step = 15
         z_range = np.arange(z_start, z_end + z_step, z_step)
 
         # Define galvo range, i.e. galvo positions at which O3 defocus stacks
@@ -787,8 +787,9 @@ def acquire(self):
         )
 
         logger.info('Starting acquisition')
+        ls_o3_refocus_time = time.time()
         for t_idx in range(self.time_settings.num_timepoints):
-            t_start = time.time()
+            timepoint_start_time = time.time()
             for p_idx in range(self.position_settings.num_positions):
                 p_label = self.position_settings.position_labels[p_idx]
 
@@ -813,7 +814,12 @@ def acquire(self):
                 # 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:
-                    self.refocus_ls_path()
+                    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
 
                 # start acquisition
                 lf_events = deepcopy(lf_cz_events)
@@ -835,7 +841,7 @@ def acquire(self):
                 self.await_cz_acq_completion()
 
             # wait for time interval between time points
-            t_wait = self.time_settings.time_interval_s - (time.time() - t_start)
+            t_wait = self.time_settings.time_interval_s - (time.time() - timepoint_start_time)
             if t_wait > 0 and t_idx < self.time_settings.num_timepoints - 1:
                 logger.info(f"Waiting {t_wait/60:.2f} minutes until the next time point")
                 time.sleep(t_wait)

From b263cff7f36262510d866388e15b846387699db2 Mon Sep 17 00:00:00 2001
From: LabelFree <ivan.ivanov@czbiohub.org>
Date: Wed, 12 Jul 2023 17:29:33 -0700
Subject: [PATCH 32/45] update kim101 calibration

---
 examples/calibrate_kim101.py                | 11 +++++++----
 mantis/acquisition/microscope_operations.py |  2 +-
 2 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/examples/calibrate_kim101.py b/examples/calibrate_kim101.py
index ea20d7aa..c79a49fa 100644
--- a/examples/calibrate_kim101.py
+++ b/examples/calibrate_kim101.py
@@ -8,14 +8,14 @@
 #%%
 import numpy as np
 from pycromanager import Core, Studio
-from mantis.acquisition.microscope_operations import setup_kim101_stage, acquire_ls_defocus_stack
+from mantis.acquisition.microscope_operations import setup_kim101_stage, acquire_ls_defocus_stack_and_display
 
 #%%
 mmc = Core()
 mmStudio = Studio()
 z_start = 0
-z_end = 200
-z_step = 25
+z_end = 150
+z_step = 15
 galvo = 'AP Galvo'
 galvo_range = [0]*5
 
@@ -28,7 +28,7 @@
 )
 
 #%%
-data = acquire_ls_defocus_stack(
+data = acquire_ls_defocus_stack_and_display(
     mmc, 
     mmStudio, 
     z_stage, 
@@ -57,3 +57,6 @@
     neg_slope.append(m)
 
 compensation_factor = np.mean(pos_slope) / np.mean(neg_slope)
+print(compensation_factor)
+
+# %%
diff --git a/mantis/acquisition/microscope_operations.py b/mantis/acquisition/microscope_operations.py
index 449e7b3a..e2c5f12b 100644
--- a/mantis/acquisition/microscope_operations.py
+++ b/mantis/acquisition/microscope_operations.py
@@ -13,7 +13,7 @@
 
 logger = logging.getLogger(__name__)
 
-KIM101_COMPENSATION_FACTOR = 1.08
+KIM101_COMPENSATION_FACTOR = 1.10
 
 
 def _try_mmc_call(mmc, mmc_call_name, *mmc_carr_args):

From 32d1857bd84098d4ffdfb51f671a6d6ef787916a Mon Sep 17 00:00:00 2001
From: LabelFree <ivan.ivanov@czbiohub.org>
Date: Wed, 12 Jul 2023 17:59:30 -0700
Subject: [PATCH 33/45] add relative O3 travel limits

---
 mantis/acquisition/acq_engine.py | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/mantis/acquisition/acq_engine.py b/mantis/acquisition/acq_engine.py
index 1af3b4a5..d4a379f8 100644
--- a/mantis/acquisition/acq_engine.py
+++ b/mantis/acquisition/acq_engine.py
@@ -626,6 +626,15 @@ def refocus_ls_path(self):
         z_step = 15
         z_range = np.arange(z_start, z_end + z_step, z_step)
 
+        # Define relative travel limits, in steps
+        z_stage = self.ls_acq.o3_stage
+        target_z_position = z_stage.true_position + z_range
+        max_z_position = 500  # O3 is allowed to travel ~10 um towards O2
+        min_z_position = -1000  # O3 is allowed to travel ~20 um away from O2
+        if np.any(target_z_position > max_z_position) or np.any(target_z_position < min_z_position):
+            logger.error('O3 relative travel limits will be exceeded. Aborting O3 refocus.')
+            return
+
         # Define galvo range, i.e. galvo positions at which O3 defocus stacks
         # are acquired, should be odd number
         galvo_scan_range = self.ls_acq.slice_settings.z_range
@@ -639,7 +648,7 @@ def refocus_ls_path(self):
         # Acquire defocus stacks at several galvo positions
         data = acquire_ls_defocus_stack(
             mmc=self.ls_acq.mmc,
-            z_stage=self.ls_acq.o3_stage,
+            z_stage=z_stage,
             z_range=z_range,
             galvo=self.ls_acq.slice_settings.z_stage_name,
             galvo_range=galvo_range,

From e5a3e3af4fa585a992369380f1d89c9c7803bad8 Mon Sep 17 00:00:00 2001
From: LabelFree <ivan.ivanov@czbiohub.org>
Date: Wed, 12 Jul 2023 18:07:28 -0700
Subject: [PATCH 34/45] logger fixes

---
 mantis/acquisition/acq_engine.py | 15 +++++++--------
 1 file changed, 7 insertions(+), 8 deletions(-)

diff --git a/mantis/acquisition/acq_engine.py b/mantis/acquisition/acq_engine.py
index d4a379f8..3789f1a9 100644
--- a/mantis/acquisition/acq_engine.py
+++ b/mantis/acquisition/acq_engine.py
@@ -302,17 +302,12 @@ def __init__(
 
         # Setup logger
         timestamp = datetime.now().strftime("%Y%m%dT%H%M%S")
-        configure_logger(
-            os.path.join(self._logs_dir, f'mantis_acquisition_log_{timestamp}.txt')
-        )
-
-        # initialize time and position settings
-        self._time_settings = TimeSettings()
-        self._position_settings = PositionSettings()
+        acq_log_path = os.path.join(self._logs_dir, f'mantis_acquisition_log_{timestamp}.txt')
+        configure_logger(acq_log_path)
 
         if self._demo_run:
             logger.info('NOTE: This is a demo run')
-        logger.debug(f'Starting mantis acquisition log at: {self._acq_dir}')
+        logger.debug(f'Starting mantis acquisition log at: {acq_log_path}')
 
         # Log conda environment
         outs, errs = log_conda_environment(
@@ -323,6 +318,10 @@ def __init__(
         else:
             logger.error(errs.decode('ascii'))
 
+        # initialize time and position settings
+        self._time_settings = TimeSettings()
+        self._position_settings = PositionSettings()
+
         # Connect to MM running LF acq
         self.lf_acq = BaseChannelSliceAcquisition(
             enabled=enable_lf_acq,

From 50ae9a1cd19e99dafb629527ecf00ede594474b9 Mon Sep 17 00:00:00 2001
From: LabelFree <ivan.ivanov@czbiohub.org>
Date: Thu, 13 Jul 2023 11:08:52 -0700
Subject: [PATCH 35/45] style

---
 mantis/acquisition/acq_engine.py            | 49 ++++++++++++---------
 mantis/acquisition/microscope_operations.py | 11 ++---
 2 files changed, 34 insertions(+), 26 deletions(-)

diff --git a/mantis/acquisition/acq_engine.py b/mantis/acquisition/acq_engine.py
index 3789f1a9..f4b85012 100644
--- a/mantis/acquisition/acq_engine.py
+++ b/mantis/acquisition/acq_engine.py
@@ -1,25 +1,24 @@
 import logging
 import os
 import time
-from typing import Iterable
 
 from copy import deepcopy
 from dataclasses import asdict
 from datetime import datetime
 from functools import partial
+from typing import Iterable
 
-import tifffile
 import nidaqmx
 import numpy as np
+import tifffile
 
 from nidaqmx.constants import Slope
 from pycromanager import Acquisition, Core, Studio, multi_d_acquisition_events, start_headless
+from waveorder.focus import focus_from_transverse_band
 
 from mantis.acquisition import microscope_operations
 from mantis.acquisition.logger import configure_logger, log_conda_environment
 
-from waveorder.focus import focus_from_transverse_band
-
 # isort: off
 from mantis.acquisition.AcquisitionSettings import (
     TimeSettings,
@@ -630,7 +629,9 @@ def refocus_ls_path(self):
         target_z_position = z_stage.true_position + z_range
         max_z_position = 500  # O3 is allowed to travel ~10 um towards O2
         min_z_position = -1000  # O3 is allowed to travel ~20 um away from O2
-        if np.any(target_z_position > max_z_position) or np.any(target_z_position < min_z_position):
+        if np.any(target_z_position > max_z_position) or np.any(
+            target_z_position < min_z_position
+        ):
             logger.error('O3 relative travel limits will be exceeded. Aborting O3 refocus.')
             return
 
@@ -639,11 +640,11 @@ def refocus_ls_path(self):
         galvo_scan_range = self.ls_acq.slice_settings.z_range
         len_galvo_scan_range = len(galvo_scan_range)
         galvo_range = [
-            galvo_scan_range[int(0.3*len_galvo_scan_range)],
-            galvo_scan_range[int(0.5*len_galvo_scan_range)],
-            galvo_scan_range[int(0.7*len_galvo_scan_range)],
+            galvo_scan_range[int(0.3 * len_galvo_scan_range)],
+            galvo_scan_range[int(0.5 * len_galvo_scan_range)],
+            galvo_scan_range[int(0.7 * len_galvo_scan_range)],
         ]
-        
+
         # Acquire defocus stacks at several galvo positions
         data = acquire_ls_defocus_stack(
             mmc=self.ls_acq.mmc,
@@ -659,7 +660,7 @@ def refocus_ls_path(self):
         timestamp = datetime.now().strftime("%Y%m%dT%H%M%S")
         tifffile.imwrite(
             os.path.join(self._logs_dir, f'ls_refocus_data_{timestamp}.ome.tif'),
-            np.expand_dims(data, -3).astype('uint16')
+            np.expand_dims(data, -3).astype('uint16'),
         )
 
         # Find in-focus slice
@@ -671,11 +672,11 @@ def refocus_ls_path(self):
             )
             focus_indices.append(idx)
         logger.debug(
-            'Stacks at galvo positions %s are in focus at slice %s', 
-            np.round(galvo_range, 3), 
-            focus_indices
+            'Stacks at galvo positions %s are in focus at slice %s',
+            np.round(galvo_range, 3),
+            focus_indices,
         )
-        
+
         # Refocus O3
         # Some focus_indices may be None, e.g. if there is no sample
         valid_focus_indices = [idx for idx in focus_indices if idx is not None]
@@ -685,12 +686,13 @@ def refocus_ls_path(self):
 
             logger.info(f'Moving O3 by {o3_displacement} steps')
             microscope_operations.set_relative_kim101_position(
-                self.ls_acq.o3_stage,
-                o3_displacement
+                self.ls_acq.o3_stage, o3_displacement
             )
         else:
-            logger.error('Could not determine the correct O3 in-focus position. O3 will not move')
-    
+            logger.error(
+                'Could not determine the correct O3 in-focus position. O3 will not move'
+            )
+
     def setup(self):
         """
         Setup the mantis acquisition. This method sets up the label-free
@@ -824,8 +826,11 @@ def acquire(self):
                 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:
+                    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
 
@@ -977,7 +982,9 @@ def acquire_ls_defocus_stack(
 
     # set camera to internal trigger
     # TODO: do this properly, context manager?
-    microscope_operations.set_property(mmc, 'Prime BSI Express', 'TriggerMode', 'Internal Trigger')
+    microscope_operations.set_property(
+        mmc, 'Prime BSI Express', 'TriggerMode', 'Internal Trigger'
+    )
 
     # acquire stack at different galvo positions
     for p_idx, p in enumerate(galvo_range):
diff --git a/mantis/acquisition/microscope_operations.py b/mantis/acquisition/microscope_operations.py
index e2c5f12b..eeb8f98f 100644
--- a/mantis/acquisition/microscope_operations.py
+++ b/mantis/acquisition/microscope_operations.py
@@ -304,7 +304,7 @@ def acquire_defocus_stack(
     z_stage,
     z_range: Iterable,
     mmStudio: Studio = None,
-    datastore = None,
+    datastore=None,
     channel_ind: int = 0,
     position_ind: int = 0,
 ):
@@ -374,6 +374,7 @@ def acquire_defocus_stack(
 
     return np.asarray(data)
 
+
 def acquire_ls_defocus_stack_and_display(
     mmc: Core,
     mmStudio: Studio,
@@ -505,10 +506,10 @@ def reset_shutter(mmc: Core, auto_shutter_state: bool, shutter_state: bool):
     shutter_device = mmc.get_shutter_device()
     if shutter_device:
         logger.debug(
-            'Resetting shutter %s to state Open:%s, Autoshutter: %s', 
-            shutter_device, 
-            shutter_state, 
-            auto_shutter_state
+            'Resetting shutter %s to state Open:%s, Autoshutter: %s',
+            shutter_device,
+            shutter_state,
+            auto_shutter_state,
         )
         mmc.set_shutter_open(shutter_state)
         mmc.set_auto_shutter(auto_shutter_state)

From 7c1e4e35c2fc7104344532273b4befd8b4c9fdc7 Mon Sep 17 00:00:00 2001
From: LabelFree <ivan.ivanov@czbiohub.org>
Date: Thu, 13 Jul 2023 11:35:39 -0700
Subject: [PATCH 36/45] add waveorder to deps and format pyproject

---
 pyproject.toml | 70 +++++++++++++++++++++++---------------------------
 1 file changed, 32 insertions(+), 38 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 632e3eab..fbe43ac5 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,60 +1,54 @@
 [build-system]
-requires = [
-  "setuptools >= 42",
-  "wheel",
-  "setuptools_scm[toml]>=3.4"
-]
+requires = ["setuptools >= 42", "setuptools_scm[toml]>=3.4", "wheel"]
 
 [project]
-name =  "mantis"
+name = "mantis"
 description = "Acquisition engine for collecting data on the mantis microscope"
 readme = "README.md"
-license = {file = "LICENSE"}
+license = { file = "LICENSE" }
 requires-python = ">=3.10, <4.0"
 
 # the dynamically determined project metadata attributes
 dynamic = ["version"]
 
 classifiers = [
-    "Development Status :: 3 - Alpha",
-    "Intended Audience :: Science/Research",
-    "Programming Language :: Python :: 3 :: Only",
-    "Programming Language :: Python :: 3.10"
+  "Development Status :: 3 - Alpha",
+  "Intended Audience :: Science/Research",
+  "Programming Language :: Python :: 3 :: Only",
+  "Programming Language :: Python :: 3.10",
 ]
 
 # list package dependencies here
 dependencies = [
-    "numpy",
-    "pydantic",
-    "pycromanager==0.25.40",
-    "nidaqmx",
-    "iohub==0.1.0.dev3",
-    "pylablib==1.4.1",
-    "scipy",
-    "slurmkit @ git+https://github.com/royerlab/slurmkit",
-    "natsort",
-    "tifffile",
-    "napari[all]; 'arm64' not in platform_machine", # with Qt5 and skimage
-    "napari; 'arm64' in platform_machine", # without Qt5 and skimage
-    "PyQt6; 'arm64' in platform_machine",
-    "matplotlib",
+  "iohub==0.1.0.dev3",
+  "matplotlib",
+  "napari; 'arm64' in platform_machine",                 # without Qt5 and skimage
+  "napari[all]; 'arm64' not in platform_machine",        # with Qt5 and skimage
+  "PyQt6; 'arm64' in platform_machine",
+  "natsort",
+  "nidaqmx",
+  "numpy",
+  "pycromanager==0.25.40",
+  "pydantic",
+  "pylablib==1.4.1",
+  "scipy",
+  "slurmkit @ git+https://github.com/royerlab/slurmkit",
+  "tifffile",
+  "waveorder",
 ]
 
 
 [project.optional-dependencies]
 # note that dev dependencies are only pinned to major versions
 dev = [
-    "black==22.3.0",
-    "flake8~=5.0",
-    "isort~=5.12",
-    "pre-commit~=2.19",
-    "pylint~=2.14",
-    "pytest~=7.1",
-]
-build = [
-    "twine",
-    "build",
+  "black==22.3.0",
+  "flake8~=5.0",
+  "isort~=5.12",
+  "pre-commit~=2.19",
+  "pylint~=2.14",
+  "pytest~=7.1",
 ]
+build = ["build", "twine"]
 
 [project.scripts]
 mantis = "mantis.cli.main:cli"
@@ -68,7 +62,7 @@ packages = ["mantis"]
 zip-safe = false
 
 [tool.setuptools.dynamic]
-version = {attr = "mantis.__version__"}
+version = { attr = "mantis.__version__" }
 
 [tool.black]
 line-length = 95
@@ -99,7 +93,7 @@ profile = "black"
 line_length = 95
 lines_between_types = 1
 default_section = "THIRDPARTY"
-no_lines_before = ["STDLIB",]
+no_lines_before = ["STDLIB"]
 ensure_newline_before_comments = true
 skip_glob = ["examples/*"]
 
@@ -107,7 +101,7 @@ skip_glob = ["examples/*"]
 # disable all conventions, refactors, warnings (C, R, W) and the following:
 # E0401: unable-to-import (since it is possible that no one environment has all required packages)
 # E1136: unsubscriptable-object (anecdotal false positives for numpy objects)
-disable = ["C", "R", "W", "unsubscriptable-object", "import-error"]
+disable = ["C", "R", "W", "import-error", "unsubscriptable-object"]
 msg-template = "{line},{column},{category},{symbol}:{msg}"
 reports = "n"
 

From 9950c5004e8e9d551c3da895cfb7b35bc275d3e5 Mon Sep 17 00:00:00 2001
From: LabelFree <ivan.ivanov@czbiohub.org>
Date: Thu, 13 Jul 2023 12:59:44 -0700
Subject: [PATCH 37/45] update kim101 compensation factor

---
 examples/acquire_defocus_stack.py           | 63 ++++++++++++---------
 examples/calibrate_kim101.py                | 12 +++-
 mantis/acquisition/microscope_operations.py |  2 +-
 3 files changed, 46 insertions(+), 31 deletions(-)

diff --git a/examples/acquire_defocus_stack.py b/examples/acquire_defocus_stack.py
index a58a70ad..4b85cd25 100644
--- a/examples/acquire_defocus_stack.py
+++ b/examples/acquire_defocus_stack.py
@@ -1,39 +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
-
-import logging
-logger = logging.getLogger()
-logger.setLevel(logging.DEBUG)
-console_handler = logging.StreamHandler()
-console_handler.setLevel(logging.DEBUG)
-console_format = logging.Formatter('%(levelname)s - %(module)s.%(funcName)s - %(message)s')
-console_handler.setFormatter(console_format)
-logger.addHandler(console_handler)
+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 = -200
-z_end = 200
-z_step = 25
-config_group = 'Channel - LS'
-config_name = 'GFP EX488 EM525-45'
+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)
 
-data = acquire_ls_defocus_stack_and_display(
-    mmc, 
-    mmStudio, 
-    z_stage, 
-    z_range,
-    galvo,
-    galvo_range,
-    config_group, 
-    config_name,
-    close_display=False,
-)
+# 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)
 
-print(data.shape)
diff --git a/examples/calibrate_kim101.py b/examples/calibrate_kim101.py
index c79a49fa..3555add5 100644
--- a/examples/calibrate_kim101.py
+++ b/examples/calibrate_kim101.py
@@ -8,18 +8,20 @@
 #%%
 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 = 150
+z_end = 105
 z_step = 15
 galvo = 'AP Galvo'
 galvo_range = [0]*5
 
-z_stage = setup_kim101_stage('74000291')
 z_range = np.hstack(
     (
         np.arange(z_start, z_end + z_step, z_step),
@@ -35,11 +37,12 @@
     z_range,
     galvo,
     galvo_range,
+    close_display = False
 )
 
 # %%
 steps_per_direction = len(z_range)//2
-intensity = data.sum(axis=(-1, -2))
+intensity = data.max(axis=(-1, -2))
 
 pos_int = intensity[:, :steps_per_direction]
 pos_z = z_range[:steps_per_direction]
@@ -60,3 +63,6 @@
 print(compensation_factor)
 
 # %%
+plt.plot(intensity.flatten())
+
+# %%
diff --git a/mantis/acquisition/microscope_operations.py b/mantis/acquisition/microscope_operations.py
index eeb8f98f..3b6a1e78 100644
--- a/mantis/acquisition/microscope_operations.py
+++ b/mantis/acquisition/microscope_operations.py
@@ -13,7 +13,7 @@
 
 logger = logging.getLogger(__name__)
 
-KIM101_COMPENSATION_FACTOR = 1.10
+KIM101_COMPENSATION_FACTOR = 1.03
 
 
 def _try_mmc_call(mmc, mmc_call_name, *mmc_carr_args):

From 04a35c68485adac433f553b079e847f7b46566a6 Mon Sep 17 00:00:00 2001
From: LabelFree <ivan.ivanov@czbiohub.org>
Date: Mon, 17 Jul 2023 16:58:40 -0700
Subject: [PATCH 38/45] add threshold and plotting to focus finding

---
 mantis/acquisition/acq_engine.py | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

diff --git a/mantis/acquisition/acq_engine.py b/mantis/acquisition/acq_engine.py
index f4b85012..e641951b 100644
--- a/mantis/acquisition/acq_engine.py
+++ b/mantis/acquisition/acq_engine.py
@@ -665,10 +665,21 @@ def refocus_ls_path(self):
 
         # Find in-focus slice
         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
+
         focus_indices = []
-        for stack in data:
+        for stack_idx, stack in enumerate(data):
             idx = focus_from_transverse_band(
-                stack, NA_det=NA_DETECTION, lambda_ill=wavelength, pixel_size=LS_PIXEL_SIZE
+                stack,
+                NA_det=NA_DETECTION,
+                lambda_ill=wavelength,
+                pixel_size=LS_PIXEL_SIZE,
+                threshold_FWHM=threshold_FWHM,
+                plot_path=os.path.join(
+                    self._logs_dir, f'ls_refocus_plot_{timestamp}_Pos{stack_idx}.png'
+                ),
             )
             focus_indices.append(idx)
         logger.debug(

From 8b580dffcd5f6a1d230e1e44a2d31c2b2e21691c Mon Sep 17 00:00:00 2001
From: Ivan Ivanov <ivan.ivanov@czbiohub.org>
Date: Mon, 17 Jul 2023 17:06:34 -0700
Subject: [PATCH 39/45] update waveorder dependency

---
 pyproject.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index fbe43ac5..611ef7d9 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -34,7 +34,7 @@ dependencies = [
   "scipy",
   "slurmkit @ git+https://github.com/royerlab/slurmkit",
   "tifffile",
-  "waveorder",
+  "waveorder @ git+https://github.com/mehta-lab/waveorder",
 ]
 
 

From f4ab651d90824f91aa01ee6aa060f10839899986 Mon Sep 17 00:00:00 2001
From: Ivan Ivanov <ivan.ivanov@czbiohub.org>
Date: Fri, 21 Jul 2023 13:24:36 -0700
Subject: [PATCH 40/45] add microscope_operations documentation

---
 mantis/acquisition/microscope_operations.py | 59 +++++++++++----------
 1 file changed, 31 insertions(+), 28 deletions(-)

diff --git a/mantis/acquisition/microscope_operations.py b/mantis/acquisition/microscope_operations.py
index 3b6a1e78..63da24e6 100644
--- a/mantis/acquisition/microscope_operations.py
+++ b/mantis/acquisition/microscope_operations.py
@@ -266,7 +266,7 @@ def setup_kim101_stage(serial_number: int, max_voltage=112, velocity=500, accele
 
 def set_relative_kim101_position(
     stage: KinesisPiezoMotor,
-    step: int,
+    distance: int,
 ):
     """Make a relative move with a KIM101 stage, compensating for different
     travel distance per step in the positive and negative directions
@@ -274,25 +274,35 @@ def set_relative_kim101_position(
     Parameters
     ----------
     stage : KinesisPiezoMotor
-        _description_
-    step : int
-        _description_
+    distance : int
+        Travel distance, in number of steps
     """
 
     # keep track of the stage actual position, in steps, not accounting for the
     # compensation factor
-    stage.true_position += step
+    stage.true_position += distance
 
-    if step < 0:
-        step *= KIM101_COMPENSATION_FACTOR
+    if distance < 0:
+        distance *= KIM101_COMPENSATION_FACTOR
 
-    stage.move_by(int(step))
+    stage.move_by(int(distance))
     stage.wait_move()
 
 
 def create_ram_datastore(
     mmStudio: Studio,
 ):
+    """Create a Micro-manager RAM datastore and associate a display window with it
+
+    Parameters
+    ----------
+    mmStudio : Studio
+
+    Returns
+    -------
+    datastore
+
+    """
     datastore = mmStudio.get_data_manager().create_ram_datastore()
     mmStudio.get_display_manager().create_display(datastore)
 
@@ -308,7 +318,7 @@ def acquire_defocus_stack(
     channel_ind: int = 0,
     position_ind: int = 0,
 ):
-    """Snap image at every z position and put image in the datastore
+    """Snap image at every z position and put image in a Micro-manager datastore
 
     Parameters
     ----------
@@ -316,8 +326,9 @@ def acquire_defocus_stack(
     z_stage : str or coPylot stage object
     z_range : Iterable
     mmStudio : Studio, optional
-        If not None, images will be added to a Micro-manager RAM datastore
-    datastore : micromanager.data.Datastore
+        If not None, images will be added to a Micro-manager RAM datastore and
+        displayed. If None, acquired images will only be returned as np.array
+    datastore : micromanager.data.Datastore, optional
         Micro-manager datastore object
     channel_ind : int, optional
         Channel index of acquired images in the Micro-manager datastore, by default 0
@@ -342,7 +353,6 @@ def acquire_defocus_stack(
     else:
         raise RuntimeError(f'Unknown z stage: {z_stage}')
 
-    # mmc.wait_for_image_synchro()
     for z_ind, rel_z in enumerate(relative_z_steps):
         # set z position
         move_z(rel_z)
@@ -378,7 +388,7 @@ def acquire_defocus_stack(
 def acquire_ls_defocus_stack_and_display(
     mmc: Core,
     mmStudio: Studio,
-    z_stage,
+    z_stage: str or KinesisPiezoMotor,
     z_range: Iterable,
     galvo: str,
     galvo_range: Iterable,
@@ -386,28 +396,21 @@ def acquire_ls_defocus_stack_and_display(
     config_name: str = None,
     close_display: bool = True,
 ):
-    """Acquire defocus stacks at different galvo positions
+    """Utility function similar to MantisAcquisition.acquire_ls_defocus_stack
+    which can acquire O3 defocus stacks at multiple galvo positions and
+    display them in a Micro-manager window
 
     Parameters
     ----------
     mmc : Core
-        _description_
     mmStudio : Studio
-        _description_
-    z_stage : _type_
-        _description_
-    z_start : float
-        _description_
-    z_end : float
-        _description_
-    z_step : float
-        _description_
+    z_stage : str or KinesisPiezoMotor
+    z_range : Iterable
+    galvo : str
+    galvo_range : Iterable
     config_group : str, optional
-        _description_, by default None
     config_name : str, optional
-        _description_, by default None
-    close_display: bool. optional
-        _description_, by default True
+    close_display : bool, optional
 
     Returns
     -------

From badfdf6b0c6fd720ce13b1c17d2d96beb2749611 Mon Sep 17 00:00:00 2001
From: Ivan Ivanov <ivan.ivanov@czbiohub.org>
Date: Fri, 21 Jul 2023 13:31:13 -0700
Subject: [PATCH 41/45] update data structure specs

---
 docs/data_structure.md | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/docs/data_structure.md b/docs/data_structure.md
index 901dea41..6ba663e2 100644
--- a/docs/data_structure.md
+++ b/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.
 

From b5e9766368bc98505782cd769eb0d13d87d0782e Mon Sep 17 00:00:00 2001
From: Ivan Ivanov <ivan.ivanov@czbiohub.org>
Date: Fri, 21 Jul 2023 13:47:02 -0700
Subject: [PATCH 42/45] make acquire_ls_defocus_stack MantisAcquisition static
 method

---
 mantis/acquisition/acq_engine.py | 154 +++++++++++++++----------------
 1 file changed, 73 insertions(+), 81 deletions(-)

diff --git a/mantis/acquisition/acq_engine.py b/mantis/acquisition/acq_engine.py
index e641951b..54293124 100644
--- a/mantis/acquisition/acq_engine.py
+++ b/mantis/acquisition/acq_engine.py
@@ -613,6 +613,76 @@ 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,
+        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
+
+        Parameters
+        ----------
+        mmc : Core
+        mmStudio : Studio
+        z_stage : str or KinesisPiezoMotor
+        z_range : Iterable
+        galvo : str
+        galvo_range : Iterable
+        config_group : str, optional
+        config_name : str, optional
+
+        Returns
+        -------
+        data : np.array
+
+        """
+        data = []
+
+        # Set config
+        if config_name is not None:
+            mmc.set_config(config_group, config_name)
+            mmc.wait_for_config(config_group, config_name)
+
+        # Open shutter
+        auto_shutter_state, shutter_state = microscope_operations.get_shutter_state(mmc)
+        microscope_operations.open_shutter(mmc)
+
+        # get galvo starting position
+        p0 = mmc.get_position(galvo)
+
+        # set camera to internal trigger
+        # TODO: do this properly, context manager?
+        microscope_operations.set_property(
+            mmc, 'Prime BSI Express', 'TriggerMode', 'Internal Trigger'
+        )
+
+        # acquire stack at different galvo positions
+        for p_idx, p in enumerate(galvo_range):
+            # set galvo position
+            mmc.set_position(galvo, p0 + p)
+
+            # acquire defocus stack
+            z_stack = microscope_operations.acquire_defocus_stack(mmc, z_stage, z_range)
+            data.append(z_stack)
+
+        # Reset camera triggering
+        microscope_operations.set_property(
+            mmc, 'Prime BSI Express', 'TriggerMode', 'Edge Trigger'
+        )
+
+        # Reset galvo
+        mmc.set_position(galvo, p0)
+
+        # Reset shutter
+        microscope_operations.reset_shutter(mmc, auto_shutter_state, shutter_state)
+
+        return np.asarray(data)
+
     def refocus_ls_path(self):
         logger.info('Running O3 refocus algorithm on light-sheet arm')
 
@@ -636,7 +706,7 @@ def refocus_ls_path(self):
             return
 
         # Define galvo range, i.e. galvo positions at which O3 defocus stacks
-        # are acquired, should be odd number
+        # are acquired, here at 30%, 50%, and 70% of galvo range. Should be odd number
         galvo_scan_range = self.ls_acq.slice_settings.z_range
         len_galvo_scan_range = len(galvo_scan_range)
         galvo_range = [
@@ -646,7 +716,7 @@ def refocus_ls_path(self):
         ]
 
         # Acquire defocus stacks at several galvo positions
-        data = acquire_ls_defocus_stack(
+        data = self.acquire_ls_defocus_stack(
             mmc=self.ls_acq.mmc,
             z_stage=z_stage,
             z_range=z_range,
@@ -656,7 +726,7 @@ def refocus_ls_path(self):
             config_name=self.ls_acq.microscope_settings.o3_refocus_config.config_name,
         )
 
-        # Save acquired stacks
+        # Save acquired stacks in logs
         timestamp = datetime.now().strftime("%Y%m%dT%H%M%S")
         tifffile.imwrite(
             os.path.join(self._logs_dir, f'ls_refocus_data_{timestamp}.ome.tif'),
@@ -938,81 +1008,3 @@ def _create_acquisition_directory(root_dir, acq_name, idx=1):
     except OSError:
         return _create_acquisition_directory(root_dir, acq_name, idx + 1)
     return acq_dir
-
-
-def acquire_ls_defocus_stack(
-    mmc: Core,
-    z_stage,
-    z_range: Iterable,
-    galvo: str,
-    galvo_range: Iterable,
-    config_group: str = None,
-    config_name: str = None,
-):
-    """Acquire defocus stacks at different galvo positions
-
-    Parameters
-    ----------
-    mmc : Core
-        _description_
-    mmStudio : Studio
-        _description_
-    z_stage : _type_
-        _description_
-    z_start : float
-        _description_
-    z_end : float
-        _description_
-    z_step : float
-        _description_
-    config_group : str, optional
-        _description_, by default None
-    config_name : str, optional
-        _description_, by default None
-    close_display: bool. optional
-        _description_, by default True
-
-    Returns
-    -------
-    data : np.array
-
-    """
-    data = []
-
-    # Set config
-    if config_name is not None:
-        mmc.set_config(config_group, config_name)
-        mmc.wait_for_config(config_group, config_name)
-
-    # Open shutter
-    auto_shutter_state, shutter_state = microscope_operations.get_shutter_state(mmc)
-    microscope_operations.open_shutter(mmc)
-
-    # get galvo starting position
-    p0 = mmc.get_position(galvo)
-
-    # set camera to internal trigger
-    # TODO: do this properly, context manager?
-    microscope_operations.set_property(
-        mmc, 'Prime BSI Express', 'TriggerMode', 'Internal Trigger'
-    )
-
-    # acquire stack at different galvo positions
-    for p_idx, p in enumerate(galvo_range):
-        # set galvo position
-        mmc.set_position(galvo, p0 + p)
-
-        # acquire defocus stack
-        z_stack = microscope_operations.acquire_defocus_stack(mmc, z_stage, z_range)
-        data.append(z_stack)
-
-    # Reset camera triggering
-    microscope_operations.set_property(mmc, 'Prime BSI Express', 'TriggerMode', 'Edge Trigger')
-
-    # Reset galvo
-    mmc.set_position(galvo, p0)
-
-    # Reset shutter
-    microscope_operations.reset_shutter(mmc, auto_shutter_state, shutter_state)
-
-    return np.asarray(data)

From 8fbfa31516a3e4932d194f8e741530bba2620d41 Mon Sep 17 00:00:00 2001
From: Ivan Ivanov <ivan.ivanov@czbiohub.org>
Date: Fri, 21 Jul 2023 13:47:11 -0700
Subject: [PATCH 43/45] update examples

---
 examples/calibrate_kim101.py                   | 7 +++++--
 examples/{find_focus.py => test_find_focus.py} | 0
 2 files changed, 5 insertions(+), 2 deletions(-)
 rename examples/{find_focus.py => test_find_focus.py} (100%)

diff --git a/examples/calibrate_kim101.py b/examples/calibrate_kim101.py
index 3555add5..1edc3e09 100644
--- a/examples/calibrate_kim101.py
+++ b/examples/calibrate_kim101.py
@@ -1,9 +1,12 @@
-# Calibration procedure 
+# 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 
+# in the positive or negative direction
+
+# This calibration procedure works alright, but could be improved
 
 #%%
 import numpy as np
diff --git a/examples/find_focus.py b/examples/test_find_focus.py
similarity index 100%
rename from examples/find_focus.py
rename to examples/test_find_focus.py

From 1f075492682f6d12c4f5fd65ce1abb3398601a17 Mon Sep 17 00:00:00 2001
From: Ivan Ivanov <ivan.ivanov@czbiohub.org>
Date: Tue, 25 Jul 2023 09:55:54 -0700
Subject: [PATCH 44/45] rename_kim101 example

---
 examples/{test_kim101_stage.py => test_kim101_copylot.py} | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)
 rename examples/{test_kim101_stage.py => test_kim101_copylot.py} (77%)

diff --git a/examples/test_kim101_stage.py b/examples/test_kim101_copylot.py
similarity index 77%
rename from examples/test_kim101_stage.py
rename to examples/test_kim101_copylot.py
index 038778bb..035dda7d 100644
--- a/examples/test_kim101_stage.py
+++ b/examples/test_kim101_copylot.py
@@ -1,4 +1,8 @@
-# %%
+# 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 15e9ee07a54936936caa51cd9b5d52a5b07e35e4 Mon Sep 17 00:00:00 2001
From: Ivan Ivanov <ivan.ivanov@czbiohub.org>
Date: Thu, 27 Jul 2023 12:38:56 -0700
Subject: [PATCH 45/45] rename z_range vars to avoid confusion

---
 mantis/acquisition/acq_engine.py | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/mantis/acquisition/acq_engine.py b/mantis/acquisition/acq_engine.py
index 54293124..cb403cf6 100644
--- a/mantis/acquisition/acq_engine.py
+++ b/mantis/acquisition/acq_engine.py
@@ -689,14 +689,14 @@ def refocus_ls_path(self):
         # 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
-        z_start = -105
-        z_end = 105
-        z_step = 15
-        z_range = np.arange(z_start, z_end + z_step, z_step)
+        o3_z_start = -105
+        o3_z_end = 105
+        o3_z_step = 15
+        o3_z_range = np.arange(o3_z_start, o3_z_end + o3_z_step, o3_z_step)
 
         # Define relative travel limits, in steps
-        z_stage = self.ls_acq.o3_stage
-        target_z_position = z_stage.true_position + z_range
+        o3_z_stage = self.ls_acq.o3_stage
+        target_z_position = o3_z_stage.true_position + o3_z_range
         max_z_position = 500  # O3 is allowed to travel ~10 um towards O2
         min_z_position = -1000  # O3 is allowed to travel ~20 um away from O2
         if np.any(target_z_position > max_z_position) or np.any(
@@ -718,8 +718,8 @@ 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=z_stage,
-            z_range=z_range,
+            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,
@@ -763,7 +763,7 @@ def refocus_ls_path(self):
         valid_focus_indices = [idx for idx in focus_indices if idx is not None]
         if valid_focus_indices:
             focus_idx = int(np.median(valid_focus_indices))
-            o3_displacement = int(z_range[focus_idx])
+            o3_displacement = int(o3_z_range[focus_idx])
 
             logger.info(f'Moving O3 by {o3_displacement} steps')
             microscope_operations.set_relative_kim101_position(