file: don't raise on missing detector channels

changelog.md

@@ -5,6 +5,7 @@
 #### Bug fixes
 
 * Fixed statistical backing returning an incorrect value.
+* Fixed bug that prevented loading an `h5` file where only a subset of the photon channels are available. This bug was introduced in Pylake `1.4.0`.
 
 ## v1.4.0 | 2024-02-28
 

lumicks/pylake/file.py

@@ -26,8 +26,10 @@ class File(Group, Force, DownsampledFD, BaselineCorrectedForce, PhotonCounts, Ph
     filename : str | os.PathLike
         The HDF5 file to open in read-only mode
 
-    rgb_to_detectors : Dict[Color, str]
-        Dictionary that maps RGB colors to a photon detector channel (either photon counts, or photon time tags)
+    rgb_to_detectors : Optional[Dict[str, str]]
+        Dictionary that maps RGB colors to a photon detector channel (either photon counts, or
+        photon time tags). Note that a channel can be left empty by providing the channel name
+        "None". Valid colors are ("Red", "Green", "Blue").
 
     Examples
     --------
@@ -48,13 +50,9 @@ class File(Group, Force, DownsampledFD, BaselineCorrectedForce, PhotonCounts, Ph
     def __init__(self, filename, *, rgb_to_detectors=None):
         import h5py
 
-        if rgb_to_detectors is None:
-            rgb_to_detectors = {"Red": "Red", "Green": "Green", "Blue": "Blue"}
-
-        self._rgb_to_detectors = rgb_to_detectors
         super().__init__(h5py.File(filename, "r"), lk_file=self)
         self._check_file_format()
-        self._check_detector_mapping()
+        self._rgb_to_detectors = self._get_detector_mapping(rgb_to_detectors)
 
     def _check_file_format(self):
         if "Bluelake version" not in self.h5.attrs:
@@ -80,21 +78,49 @@ def _check_file_format(self):
             "Point Scan": ("point_scans", PointScan),
         }
 
-    def _check_detector_mapping(self):
-        detectors = set()
-        if "Photon time tags" in self.h5:
-            detectors = set(self.h5["Photon time tags"])
-        elif "Photon count" in self.h5:
-            detectors = set(self.h5["Photon count"])
+    def _get_detector_mapping(self, rgb_to_detectors=None):
+        """Returns the detector mapping to be used.
 
-        # Only check if detector data was exported
-        if not detectors:
-            return
+        Parameters
+        ----------
+        rgb_to_detectors : Optional[Dict[str, str]]
+            Dictionary that maps RGB colors to a photon detector channel (either photon counts, or
+            photon time tags). Note that a channel can be left empty by providing the channel name
+            "None".
+        """
 
-        if not_found := set(self._rgb_to_detectors.values()) - detectors:
-            raise Exception(
-                f"Invalid RGB to detector mapping: {not_found} photon count channel(s) are missing, images cant't be reconstructed. Available detectors: {detectors}"
-            )
+        def check_custom_detector_mapping():
+            for key in rgb_to_detectors.keys():
+                if key not in ("Red", "Green", "Blue"):
+                    raise ValueError(
+                        f'Invalid color mapping ({key}). Valid colors are "Red", "Green" or "Blue"'
+                    )
+
+            detectors = set()
+            if "Photon time tags" in self.h5:
+                detectors = set(self.h5["Photon time tags"])
+            elif "Photon count" in self.h5:
+                detectors = set(self.h5["Photon count"])
+
+            # Only check if detector data was exported
+            if not detectors:
+                return
+
+            # "None" indicates that the user doesn't want to plot data for that particular channel.
+            if not_found := (set(rgb_to_detectors.values()) - {"None"}) - detectors:
+                warnings.warn(
+                    RuntimeWarning(
+                        f"Invalid RGB to detector mapping: {not_found} photon count channel(s) are "
+                        f"missing. Those channels will be blank in images. Available detectors: "
+                        f"{detectors}"
+                    )
+                )
+
+        if rgb_to_detectors:
+            check_custom_detector_mapping()
+            return rgb_to_detectors
+        else:
+            return {"Red": "Red", "Green": "Green", "Blue": "Blue"}
 
     @classmethod
     def from_h5py(cls, h5py_file, *, rgb_to_detectors=None):
@@ -103,10 +129,7 @@ def from_h5py(cls, h5py_file, *, rgb_to_detectors=None):
         new_file.h5 = h5py_file
         new_file._lk_file = new_file
         new_file._check_file_format()
-        if rgb_to_detectors is None:
-            rgb_to_detectors = {"Red": "Red", "Green": "Green", "Blue": "Blue"}
-        new_file._rgb_to_detectors = rgb_to_detectors
-        new_file._check_detector_mapping()
+        new_file._rgb_to_detectors = new_file._get_detector_mapping(rgb_to_detectors)
         return new_file
 
     @property

lumicks/pylake/tests/test_file/conftest.py

@@ -223,3 +223,18 @@ def h5_custom_detectors(tmpdir_factory, request):
     mock_file.make_continuous_channel("Photon count", "Detector 3", np.int64(20e9), freq, counts)
     mock_file.make_continuous_channel("Info wave", "Info wave", np.int64(20e9), freq, infowave)
     return mock_file.file
+
+
+@pytest.fixture(scope="module", params=[MockDataFile_v2])
+def h5_two_colors(tmpdir_factory, request):
+    mock_class = request.param
+
+    tmpdir = tmpdir_factory.mktemp("pylake")
+    mock_file = mock_class(tmpdir.join("%s.h5" % mock_class.__class__.__name__))
+    mock_file.write_metadata()
+
+    freq = 1e9 / 16
+    mock_file.make_continuous_channel("Photon count", "Red", np.int64(20e9), freq, counts)
+    mock_file.make_continuous_channel("Photon count", "Blue", np.int64(20e9), freq, counts)
+    mock_file.make_continuous_channel("Info wave", "Info wave", np.int64(20e9), freq, infowave)
+    return mock_file.file

lumicks/pylake/tests/test_file/test_file.py

@@ -471,11 +471,64 @@ def cropped_h5(name, crop=None):
             )
 
 
-def test_detector_mapping(h5_custom_detectors):
-    f = pylake.File.from_h5py(
+def test_detector_mapping(h5_custom_detectors, h5_two_colors):
+    custom_correct = pylake.File.from_h5py(
         h5_custom_detectors,
         rgb_to_detectors={"Red": "Detector 1", "Green": "Detector 2", "Blue": "Detector 3"},
     )
-    with pytest.raises(Exception):
-        f = pylake.File.from_h5py(h5_custom_detectors, match="Invalid RGB to detector mapping")
-    assert np.any(f.red_photon_count.data)
+    assert custom_correct.red_photon_count.data.size
+    assert custom_correct.green_photon_count.data.size
+    assert custom_correct.blue_photon_count.data.size
+
+    # None of these exist, since this is a custom photon count file
+    with pytest.warns(RuntimeWarning, match="Invalid RGB to detector mapping"):
+        custom_empty = pylake.File.from_h5py(
+            h5_custom_detectors, rgb_to_detectors={"Red": "Red", "Green": "Green", "Blue": "Blue"}
+        )
+        assert not custom_empty.red_photon_count.data.size
+        assert not custom_empty.green_photon_count.data.size
+        assert not custom_empty.blue_photon_count.data.size
+
+    # Green photon count channel doesn't exist, so we issue a warning about this
+    with pytest.warns(RuntimeWarning, match="Invalid RGB to detector mapping"):
+        two_color = pylake.File.from_h5py(
+            h5_two_colors, rgb_to_detectors={"Red": "Red", "Green": "Green", "Blue": "Blue"}
+        )
+        assert two_color.red_photon_count.data.size
+        assert not two_color.green_photon_count.data.size
+        assert two_color.blue_photon_count.data.size
+
+    # Green photon count channel doesn't exist, _but_ we are using defaults, so we do not explicitly
+    # issue a warning (this is the old reference behavior). People may not even want to show
+    # images, so it's weird to bother them with a warning.
+    two_color_default = pylake.File.from_h5py(h5_two_colors)
+    assert two_color_default.red_photon_count.data.size
+    assert not two_color_default.green_photon_count.data.size
+    assert two_color_default.blue_photon_count.data.size
+
+    # Valid mapping
+    double_red = pylake.File.from_h5py(
+        h5_two_colors, rgb_to_detectors={"Red": "Red", "Green": "Red", "Blue": "Blue"}
+    )
+    np.testing.assert_allclose(double_red.green_photon_count.data, two_color.red_photon_count.data)
+    assert double_red.blue_photon_count.data.size
+
+    # Explicitly setting `"None"` doesn't issue a warning
+    no_green = pylake.File.from_h5py(
+        h5_two_colors, rgb_to_detectors={"Red": "Red", "Green": "None", "Blue": "Blue"}
+    )
+    assert no_green.red_photon_count.data.size
+    assert not no_green.green_photon_count.data.size
+    assert no_green.blue_photon_count.data.size
+
+
+@pytest.mark.parametrize(
+    "mapping",
+    [
+        {"red": "Red", "green": "Red", "blue": "Blue"},  # Incorrect capitalization
+        {"Red": "Red", "Magenta": "Red", "Blue": "Blue"},  # Wrong name
+    ],
+)
+def test_detector_mapping_invalid_color(h5_two_colors, mapping):
+    with pytest.raises(ValueError, match="Invalid color mapping"):
+        pylake.File.from_h5py(h5_two_colors, rgb_to_detectors=mapping)