kymo: add PositionUnit

lumicks/pylake/kymo.py

@@ -1,5 +1,7 @@
 import warnings
 from copy import copy
+from enum import Enum
+from collections import namedtuple
 from dataclasses import dataclass
 
 import numpy as np
@@ -99,7 +101,7 @@ def __init__(self, name, file, start, stop, metadata, position_offset=0, calibra
             else (
                 PositionCalibration()
                 if self.pixelsize_um[0] is None
-                else PositionCalibration("um", self.pixelsize_um[0], r"μm")
+                else PositionCalibration(PositionUnit.um, self.pixelsize_um[0])
             )
         )
 
@@ -785,7 +787,7 @@ def estimate_bead_edges(
         RuntimeError
             When the algorithm fails to locate two edges during the bead finding stage.
         """
-        if self._calibration.unit != "um":
+        if self._calibration.unit != PositionUnit.um:
             raise RuntimeError(
                 f"This kymograph is not calibrated in um but in {self._calibration.unit}. "
                 f"Please make sure the kymograph is calibrated to microns before using this "
@@ -1004,11 +1006,11 @@ def calibrate_to_kbp(self, length_kbp):
         length : float
             length of the kymo in kilobase pairs
         """
-        if self._calibration.unit == "kbp":
+        if self._calibration.unit == PositionUnit.kbp:
             raise RuntimeError("kymo is already calibrated in base pairs.")
 
         result = copy(self)
-        result._calibration = PositionCalibration("kbp", length_kbp / self._num_pixels[0], "kbp")
+        result._calibration = PositionCalibration(PositionUnit.kbp, length_kbp / self._num_pixels[0])
         result._image_factory = self._image_factory
         result._timestamp_factory = self._timestamp_factory
         result._line_time_factory = self._line_time_factory
@@ -1082,17 +1084,39 @@ def __bool__(self):
         return False
 
 
+UnitInfo = namedtuple("UnitInfo", ["name", "label"])
+
+
+class PositionUnit(Enum):
+    um = UnitInfo(name="um", label=r"μm")
+    kbp = UnitInfo(name="kbp", label="kbp")
+    pixel = UnitInfo(name="pixel", label="pixels")
+
+    def __str__(self):
+        return self.value.name
+
+    def __hash__(self):
+        return hash(self.value)
+
+    @property
+    def label(self):
+        return self.value.label
+
+
 @dataclass(frozen=True)
 class PositionCalibration:
-    unit: str = "pixel"
+    unit: PositionUnit = PositionUnit.pixel
     value: float = 1.0
-    unit_label: str = "pixels"
+
+    @property
+    def unit_label(self):
+        return self.unit.label
 
     def downsample(self, factor):
         return (
             self
-            if self.unit == "pixel"
-            else PositionCalibration(self.unit, self.value * factor, self.unit_label)
+            if self.unit == PositionUnit.pixel
+            else PositionCalibration(self.unit, self.value * factor)
         )
 
 

lumicks/pylake/kymotracker/kymotrack.py

@@ -106,7 +106,7 @@ def export_kymotrackgroup_to_csv(
     )
 
     time_units = "seconds"
-    position_units = kymotrack_group._calibration_info["unit"]
+    position_units = kymotrack_group._calibration_info.unit
 
     idx = np.hstack([np.full(len(track), idx) for idx, track in enumerate(kymotrack_group)])
     coords_idx = np.hstack([track.coordinate_idx for track in kymotrack_group])
@@ -1213,7 +1213,7 @@ def _validate_single_linetime_pixelsize(self):
             if len(pixel_sizes) == 1
             else (
                 "All source kymographs must have the same pixel sizes, "
-                f"got {sorted(pixel_sizes)} {self._calibration_info['unit']}."
+                f"got {sorted(pixel_sizes)} {self._calibration_info.unit}."
             )
         )
 
@@ -1281,7 +1281,7 @@ def _channel(self):
     def _calibration_info(self):
         try:
             kymo = self._kymos[0]
-            return {"unit": kymo._calibration.unit, "unit_label": kymo._calibration.unit_label}
+            return kymo._calibration
         except IndexError:
             raise RuntimeError("No kymo associated with this empty group (no tracks available)")
 
@@ -1470,7 +1470,7 @@ def plot(self, *, show_outline=True, show_labels=True, axes=None, **kwargs):
             track.plot(show_outline=show_outline, show_labels=False, axes=ax, **kwargs)
 
         if show_labels:
-            ax.set_ylabel(f"position ({self._calibration_info['unit_label']})")
+            ax.set_ylabel(f"position ({self._calibration_info.unit_label})")
             ax.set_xlabel("time (s)")
 
     def _tracks_in_frame(self, frame_idx):
@@ -1899,7 +1899,7 @@ def plot_binding_histogram(self, kind, bins=10, **kwargs):
         widths = np.diff(edges)
         plt.bar(edges[:-1], counts, width=widths, align="edge", **kwargs)
         plt.ylabel("Counts")
-        plt.xlabel(f"Position ({self._calibration_info['unit_label']})")
+        plt.xlabel(f"Position ({self._calibration_info.unit_label})")
 
     def _histogram_binding_profile(self, n_time_bins, bandwidth, n_position_points, roi=None):
         """Calculate a Kernel Density Estimate (KDE) of binding density along the tether for time bins.
@@ -2146,5 +2146,6 @@ def ensemble_msd(self, max_lag=None, min_count=2) -> EnsembleMSD:
             line_msds=track_msds,
             time_step=self._kymos[0].line_time_seconds,
             min_count=min_count,
-            **self._calibration_info,
+            unit=self._calibration_info.unit,
+            unit_label=self._calibration_info.unit_label,
         )

lumicks/pylake/kymotracker/kymotracker.py

@@ -3,6 +3,7 @@
 
 import numpy as np
 
+from ..kymo import PositionUnit
 from .kymotrack import KymoTrack, KymoTrackGroup
 from .detail.peakfinding import find_kymograph_peaks, refine_peak_based_on_moment
 from .detail.gaussian_mle import gaussian_mle_1d, overlapping_pixels
@@ -19,7 +20,11 @@
 ]
 
 
-_default_track_widths = {"um": 0.35, "kbp": 0.35 / 0.34, "pixel": 4}
+_default_track_widths = {
+    PositionUnit.um: 0.35,
+    PositionUnit.kbp: 0.35 / 0.34,
+    PositionUnit.pixel: 4,
+}
 
 
 def _to_pixel_rect(rect, pixelsize, line_time_seconds):

lumicks/pylake/kymotracker/tests/test_kymotrackgroup_sources.py

@@ -2,7 +2,7 @@
 
 import pytest
 
-from lumicks.pylake.kymo import _kymo_from_array
+from lumicks.pylake.kymo import PositionUnit, _kymo_from_array
 from lumicks.pylake.kymotracker.kymotrack import *
 
 
@@ -81,7 +81,7 @@ def test_constructor(kymos, coordinates):
     tracks = KymoTrackGroup(raw_tracks)
     assert len(tracks) == 4
     assert tracks._kymos == (kymo,)
-    check_attributes(tracks, [10e-4], [0.05], ["um"])
+    check_attributes(tracks, [10e-4], [0.05], [PositionUnit.um])
 
     # construct with duplicate track
     with pytest.raises(
@@ -117,13 +117,13 @@ def test_extend_single_source(kymos, coordinates):
     tracks = tracks1 + raw_tracks[2]
     assert len(tracks) == 3
     assert tracks._kymos == (kymo,)
-    check_attributes(tracks, [10e-4], [0.05], ["um"])
+    check_attributes(tracks, [10e-4], [0.05], [PositionUnit.um])
 
     # add group
     tracks = tracks1 + tracks2
     assert len(tracks) == 4
     assert tracks._kymos == (kymo,)
-    check_attributes(tracks, [10e-4], [0.05], ["um"])
+    check_attributes(tracks, [10e-4], [0.05], [PositionUnit.um])
 
     # extend with single source, different channels
     with pytest.raises(
@@ -154,18 +154,18 @@ def test_extend_empty(kymos, coordinates):
     tracks = empty + tracks2[0]
     assert len(tracks) == 1
     assert tracks._kymos == (kymo,)
-    check_attributes(tracks, [10e-4], [0.05], ["um"])
+    check_attributes(tracks, [10e-4], [0.05], [PositionUnit.um])
 
     # add group
     tracks = empty + tracks2
     assert len(tracks) == 4
     assert tracks._kymos == (kymo,)
-    check_attributes(tracks, [10e-4], [0.05], ["um"])
+    check_attributes(tracks, [10e-4], [0.05], [PositionUnit.um])
 
     tracks = tracks2 + empty
     assert len(tracks) == 4
     assert tracks._kymos == (kymo,)
-    check_attributes(tracks, [10e-4], [0.05], ["um"])
+    check_attributes(tracks, [10e-4], [0.05], [PositionUnit.um])
 
 
 def test_different_sources_same_attributes(kymos, coordinates):
@@ -190,7 +190,7 @@ def test_different_sources_same_attributes(kymos, coordinates):
     tracks = tracks1[:2] + tracks2[2:]
     assert len(tracks) == 4
     assert tracks._kymos == (kymo1, kymo2)
-    check_attributes(tracks, [10e-4], [0.05], ["um"])
+    check_attributes(tracks, [10e-4], [0.05], [PositionUnit.um])
 
 
 def test_different_sources_different_attributes(kymos, coordinates):
@@ -214,7 +214,7 @@ def make_tracks(kymo):
     tracks = tracks1 + tracks2
     assert len(tracks) == 8
     assert tracks._kymos == (kymos[0], kymos[2])
-    check_attributes(tracks, [10e-4, 10e-3], [0.05], ["um"])
+    check_attributes(tracks, [10e-4, 10e-3], [0.05], [PositionUnit.um])
 
     # different pixel sizes
     tracks1 = make_tracks(kymos[0])
@@ -223,7 +223,7 @@ def make_tracks(kymo):
     tracks = tracks1 + tracks2
     assert len(tracks) == 8
     assert tracks._kymos == (kymos[0], kymos[1])
-    check_attributes(tracks, [10e-4], [0.05, 0.1], ["um"])
+    check_attributes(tracks, [10e-4], [0.05, 0.1], [PositionUnit.um])
 
     # different line times and pixel sizes
     tracks1 = make_tracks(kymos[0])
@@ -232,7 +232,7 @@ def make_tracks(kymo):
     tracks = tracks1 + tracks2
     assert len(tracks) == 8
     assert tracks._kymos == (kymos[0], kymos[-1])
-    check_attributes(tracks, [10e-4, 10e-3], [0.05, 0.1], ["um"])
+    check_attributes(tracks, [10e-4, 10e-3], [0.05, 0.1], [PositionUnit.um])
 
     # extend with different calibrations
     with pytest.raises(

lumicks/pylake/nb_widgets/tests/test_image_editing.py

@@ -5,6 +5,7 @@
 import matplotlib.pyplot as plt
 
 from lumicks.pylake import ImageStack
+from lumicks.pylake.kymo import PositionUnit
 from lumicks.pylake.detail.widefield import TiffStack
 from lumicks.pylake.nb_widgets.image_editing import KymoEditorWidget, ImageEditorWidget
 from lumicks.pylake.tests.data.mock_widefield import MockTiffFile, make_alignment_image_data
@@ -116,7 +117,7 @@ def test_kymo_cropping_clicks(kymograph, region_select):
     np.testing.assert_equal(ax.position_limits, (1, 7))
 
     new_kymo = w.kymo
-    assert new_kymo._calibration.unit == "um"
+    assert new_kymo._calibration.unit == PositionUnit.um
     np.testing.assert_equal(new_kymo.get_image("red").shape, (16, 12))
 
     # with calibration
@@ -129,5 +130,5 @@ def test_kymo_cropping_clicks(kymograph, region_select):
     np.testing.assert_equal(ax.position_limits, (1, 7))
 
     new_kymo = w.kymo
-    assert new_kymo._calibration.unit == "kbp"
+    assert new_kymo._calibration.unit == PositionUnit.kbp
     np.testing.assert_equal(new_kymo.get_image("red").shape, (16, 12))

lumicks/pylake/tests/test_imaging_confocal/test_kymo_from_array.py

@@ -4,7 +4,7 @@
 import pytest
 import matplotlib.pyplot as plt
 
-from lumicks.pylake.kymo import _kymo_from_array
+from lumicks.pylake.kymo import PositionUnit, _kymo_from_array
 
 timestamp_err_msg = (
     "Per-pixel timestamps are not implemented. "
@@ -133,7 +133,7 @@ def test_from_array_no_pixelsize(test_kymo):
 
     assert arr_kymo.pixelsize_um == [None]
     assert arr_kymo.pixelsize == [1.0]
-    assert arr_kymo._calibration.unit == "pixel"
+    assert arr_kymo._calibration.unit == PositionUnit.pixel
 
     assert arr_kymo._metadata.center_point_um == {key: None for key in ("x", "y", "z")}
     assert arr_kymo._metadata.num_frames == 0

lumicks/pylake/tests/test_imaging_confocal/test_kymo_transforms.py

@@ -1,6 +1,8 @@
 import numpy as np
 import pytest
 
+from lumicks.pylake.kymo import PositionUnit
+
 
 def test_calibrate_to_kbp(test_kymo):
     kymo, ref = test_kymo
@@ -10,11 +12,11 @@ def test_calibrate_to_kbp(test_kymo):
     kymo_bp = kymo.calibrate_to_kbp(length_kbp)
 
     # test that default calibration is in microns
-    assert kymo._calibration.unit == "um"
+    assert kymo._calibration.unit == PositionUnit.um
     assert kymo._calibration.value == 0.1
 
     # test that calibration is stored as kilobase-pairs
-    assert kymo_bp._calibration.unit == "kbp"
+    assert kymo_bp._calibration.unit == PositionUnit.kbp
     np.testing.assert_allclose(kymo_bp._calibration.value, length_kbp / n_pixels)
 
     # test conversion from microns to calibration units
@@ -100,3 +102,18 @@ def test_flip_kymo(test_kymo, crop):
             kymo_flipped.flip().get_image(channel=channel),
             kymo.get_image(channel=channel),
         )
+
+
+def test_position_unit():
+    assert PositionUnit.um.label == r"μm"
+    assert PositionUnit.kbp.label == "kbp"
+    assert PositionUnit.pixel.label == "pixels"
+
+    assert str(PositionUnit.um) == "um"
+    assert str(PositionUnit.kbp) == "kbp"
+    assert str(PositionUnit.pixel) == "pixel"
+
+    assert {PositionUnit.um, PositionUnit.um, PositionUnit.kbp} == {
+        PositionUnit.um,
+        PositionUnit.kbp,
+    }