kymo: allow custom titles in plot_with_channels

changelog.md

@@ -6,6 +6,7 @@
 
 * Added parameter `allow_overwrite` to [`lk.download_from_doi()`](https://lumicks-pylake.readthedocs.io/en/latest/_api/lumicks.pylake.download_from_doi.html#lumicks.pylake.download_from_doi) to allow re-downloading only those files where the checksum does not match.
 * Added force calibration information to channels accessed directly via the square bracket notation (e.g. `file["Force HF"]["Force 1x"].calibration`).
+* Added parameter `titles` to customize title of each subplot in `[`Kymo.plot_with_channels()`](https://lumicks-pylake.readthedocs.io/en/latest/_api/lumicks.pylake.kymo.Kymo.html#lumicks.pylake.kymo.Kymo.plot_with_channels).
 
 ## v1.5.1 | 2024-06-03
 

docs/tutorial/kymographs.rst

@@ -280,7 +280,7 @@ We achieved this by passing `np.sum` to the `reduce` parameter of :func:`~lumick
 This results in summing the photon counts rather than taking their average.
 The argument `title_vertical=True` places the channel names along the y-axis instead of the axis title allowing a slightly more compact plot.
 
-Note that the plot can be further customized by specifying custom `labels`, `colors` and a `scale_bar`::
+Note that the plot can be further customized by specifying custom `labels`, `titles`, `colors` and a `scale_bar`::
 
     kymo.plot_with_channels(
         [
@@ -294,6 +294,7 @@ Note that the plot can be further customized by specifying custom `labels`, `col
         scale_bar=lk.ScaleBar(10.0, 5.0),
         colors=[[1.0, 0.2, 0.5], "green"],
         labels=["My force", "My photons"],
+        titles=["", "Step-wise forces", "Line-averaged photons"],
     )
 
 .. image:: ./figures/kymographs/kymo_plot_with_channels_customized.png

lumicks/pylake/kymo.py

@@ -426,6 +426,7 @@ def plot_with_channels(
         labels=None,
         colors=None,
         scale_bar=None,
+        titles=None,
         **kwargs,
     ):
         """Plot kymo with channel data.
@@ -450,6 +451,9 @@ def plot_with_channels(
             Forwarded to color argument of :func:`matplotlib.pyplot.plot()`.
         scale_bar : lk.ScaleBar, optional
             Scale bar to add to the kymograph.
+        titles : str | List[str]
+            List of custom titles for each subplot. Note that this requires providing a title for
+            each subplot, including the kymograph.
         **kwargs
             Forwarded to :meth:`Slice.plot() <lumicks.pylake.channel.Slice.plot()>`.
 
@@ -480,11 +484,11 @@ def set_aspect_ratio(axis, ar):
             """This function forces a specific aspect ratio, can be useful when aligning figures"""
             axis.set_aspect(ar * np.abs(np.diff(axis.get_xlim())[0] / np.diff(axis.get_ylim()))[0])
 
-        def check_length(items, item_type):
-            if len(items) != len(channels):
+        def check_length(items, item_type, which, expected_length):
+            if len(items) != expected_length:
                 raise ValueError(
                     f"When a list of {item_type} is provided, it needs to have the same length as "
-                    f"the number of channels provided. Expected {len(channels)}, got: "
+                    f"the number of {which}. Expected {expected_length}, got: "
                     f"{len(items)}."
                 )
 
@@ -501,11 +505,15 @@ def check_length(items, item_type):
 
         if labels:
             labels = [labels] if isinstance(labels, str) else labels
-            check_length(labels, "labels")
+            check_length(labels, "labels", "channels provided", len(channels))
 
         if colors:
             colors = [colors] if is_color_like(colors) else colors
-            check_length(colors, "colors")
+            check_length(colors, "colors", "channels provided", len(channels))
+
+        if titles:
+            titles = [titles] if isinstance(titles, str) else titles
+            check_length(titles, "titles", "subplots", len(channels) + 1)
 
         _, axes = plt.subplots(len(channels) + 1, 1, sharex="all")
 
@@ -540,6 +548,10 @@ def check_length(items, item_type):
         for ax in axes:
             set_aspect_ratio(ax, aspect_ratio)
 
+        if titles:
+            for ax, title in zip(axes, titles):
+                ax.set_title(title)
+
         for ax in axes[:-1]:
             ax.set_xlabel(None)
 

lumicks/pylake/tests/test_imaging_confocal/test_kymo_plotting.py

@@ -86,6 +86,7 @@ def plot_channels():
             force_over_kymolines,
             color_channel="red",
             labels="single",
+            titles=["kymo_title", "plot_title"],
             colors=[0, 0, 1],
             scale_bar=lk.ScaleBar(5.0, 5.0),
         )
@@ -95,6 +96,7 @@ def plot_with_multiple_channels():
             [force_over_kymolines, force_over_kymolines],
             color_channel="red",
             labels=["f2x", "f2x"],
+            titles=["kymo_title", "plot1", "plot2"],
             colors=["red", [1.0, 0.0, 0.1]],
         )
 
@@ -126,9 +128,9 @@ def test_plotting_labels(kymo_h5_file):
     f = lk.File.from_h5py(kymo_h5_file)
     kymo = f.kymos["tester"]
 
-    def check_axes(labels, titles):
+    def check_axes(labels, titles, first=1):
         axes = plt.gcf().get_axes()
-        for label, title, ax in zip(labels, titles, axes[1:]):
+        for label, title, ax in zip(labels, titles, axes[first:]):
             assert ax.get_title() == title
             assert ax.get_ylabel() == label
 
@@ -152,6 +154,9 @@ def check_axes(labels, titles):
     kymo.plot_with_channels([f.force2x, f["Photon count"]["Red"]], color_channel="red")
     check_axes(["Force (pN)", "y"], ["Force HF/Force 2x", "Photon count/Red"])
 
+    kymo.plot_with_channels(f.force2x, color_channel="red", titles=["hello", "this is force"])
+    check_axes([r"position (μm)", "Force (pN)"], ["hello", "this is force"], first=0)
+
 
 def test_plotting_with_channels_bad_args(kymo_h5_file):
     f = lk.File.from_h5py(kymo_h5_file)