diff --git a/changelog.md b/changelog.md index 41fd7d02c..0bc1b14ee 100644 --- a/changelog.md +++ b/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 diff --git a/lumicks/pylake/file.py b/lumicks/pylake/file.py index 58e9ecf3e..604ef978e 100644 --- a/lumicks/pylake/file.py +++ b/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 diff --git a/lumicks/pylake/tests/test_file/conftest.py b/lumicks/pylake/tests/test_file/conftest.py index 008031cc6..0fc22cd8e 100644 --- a/lumicks/pylake/tests/test_file/conftest.py +++ b/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 diff --git a/lumicks/pylake/tests/test_file/test_file.py b/lumicks/pylake/tests/test_file/test_file.py index 19466c902..0cf12b300 100644 --- a/lumicks/pylake/tests/test_file/test_file.py +++ b/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)