diff --git a/src/murfey/workflows/clem/register_preprocessing_results.py b/src/murfey/workflows/clem/register_preprocessing_results.py index 1d3acbe01..90a7113ad 100644 --- a/src/murfey/workflows/clem/register_preprocessing_results.py +++ b/src/murfey/workflows/clem/register_preprocessing_results.py @@ -14,7 +14,7 @@ from functools import cached_property from importlib.metadata import entry_points from pathlib import Path -from typing import Literal, Optional +from typing import Literal, Optional, TypeAlias from pydantic import BaseModel, computed_field from sqlmodel import Session, select @@ -29,24 +29,24 @@ logger = logging.getLogger("murfey.workflows.clem.register_preprocessing_results") +ColorChannels: TypeAlias = Literal[ + "gray", "red", "green", "blue", "cyan", "magenta", "yellow" +] + +DENOISING_MODES = ("_ICC", "_Lng_LVCC", "_Lng_SVCC") + class CLEMPreprocessingResult(BaseModel): series_name: str number_of_members: int is_stack: bool is_montage: bool - output_files: dict[ - Literal["gray", "red", "green", "blue", "cyan", "magenta", "yellow"], Path - ] - thumbnails: dict[ - Literal["gray", "red", "green", "blue", "cyan", "magenta", "yellow"], Path - ] = {} + output_files: dict[ColorChannels, Path] + thumbnails: dict[ColorChannels, Path] = {} thumbnail_size: Optional[tuple[int, int]] = None # height, width metadata: Path parent_lif: Optional[Path] = None - parent_tiffs: dict[ - Literal["gray", "red", "green", "blue", "cyan", "magenta", "yellow"], list[Path] - ] = {} + parent_tiffs: dict[ColorChannels, list[Path]] = {} pixels_x: int pixels_y: int units: str @@ -59,23 +59,35 @@ class CLEMPreprocessingResult(BaseModel): @cached_property def is_denoised(self) -> bool: """ - The "_Lng_LVCC" and "_Lng_SVCC" suffixes appended to a CLEM dataset's position - name indicate that it's a denoised image set of the same position. They should - override or supersede the original ones if they're present + The "_ICC", "_Lng_LVCC", and "_Lng_SVCC" suffixes appended to a CLEM dataset's + position name indicate that it's a denoised image set of the same position. + They should override or supersede the original ones if they're present. """ - return any( - pattern in self.series_name for pattern in ("_Lng_LVCC", "_Lng_SVCC") - ) + return any(self.series_name.endswith(pattern) for pattern in DENOISING_MODES) + + # Vallid Pydantic decorator not supported by MyPy + @computed_field # type: ignore + @cached_property + def denoising_mode(self) -> str | None: + """ + Store the denoising mode used as an attribute + """ + for pattern in DENOISING_MODES: + if self.series_name.endswith(pattern): + return pattern[1:] + return None # Valid Pydantic decorator not supported by MyPy @computed_field # type: ignore @cached_property def site_name(self) -> str: """ - Extract just the name of the site by removing the "_Lng_LVCC" suffix from + Extract just the name of the site by removing the denoising mode suffix from the series name. """ - return self.series_name.replace("_Lng_LVCC", "").replace("_Lng_SVCC", "") + if self.denoising_mode is not None: + return self.series_name[: -(len(self.denoising_mode) + 1)] + return self.series_name # Valid Pydantic decorator not supported by MyPy @computed_field # type: ignore diff --git a/tests/workflows/clem/test_register_preprocessing_results.py b/tests/workflows/clem/test_register_preprocessing_results.py index 752aea1be..b17d3d1b5 100644 --- a/tests/workflows/clem/test_register_preprocessing_results.py +++ b/tests/workflows/clem/test_register_preprocessing_results.py @@ -37,6 +37,7 @@ def generate_preprocessing_messages( rsync_basepath: Path, session_id: int, colors: list[str], + denoising_suffix: str, ): # Make directory to where data for current grid is stored visit_dir = rsync_basepath / "2020" / visit_name @@ -86,7 +87,7 @@ def generate_preprocessing_messages( datasets.extend( [ ( - grid_dir / "TileScan 1" / f"Position {n + 1}_Lng_LVCC", + grid_dir / "TileScan 1" / f"Position {n + 1}{denoising_suffix}", True, False, (2048, 2048), @@ -286,6 +287,7 @@ def test_run( rsync_basepath=rsync_basepath, session_id=ExampleVisit.murfey_session_id, colors=colors, + denoising_suffix="_Lng_LVCC", ) for message in preprocessing_messages: result = run( @@ -305,13 +307,13 @@ def test_run( @pytest.mark.parametrize( "test_params", ( - # Reverse list order? | Colors - (False, ["gray"]), - (True, ["gray"]), - (False, ["red", "green", "blue"]), - (True, ["cyan", "magenta", "yellow"]), - (False, ["gray", "red", "green", "blue"]), - (True, ["gray", "cyan", "magenta", "yellow"]), + # Reverse list order? | Colors | Denoising suffix + (False, ["gray"], "_Lng_LVCC"), + (True, ["gray"], "_Lng_SVCC"), + (False, ["red", "green", "blue"], "_ICC"), + (True, ["cyan", "magenta", "yellow"], "_Lng_LVCC"), + (False, ["gray", "red", "green", "blue"], "_Lng_SVCC"), + (True, ["gray", "cyan", "magenta", "yellow"], "_ICC"), ), ) def test_run_with_db( @@ -320,10 +322,10 @@ def test_run_with_db( mock_ispyb_credentials, murfey_db_session: SQLModelSession, ispyb_db_session: SQLAlchemySession, - test_params: tuple[bool, list[str]], + test_params: tuple[bool, list[str], str], ): # Unpack test params - (shuffle_message, colors) = test_params + (shuffle_message, colors, denoising_suffix) = test_params # Create a session to insert for this test murfey_session: MurfeyDB.Session = get_or_create_db_entry( @@ -381,6 +383,7 @@ def test_run_with_db( rsync_basepath=rsync_basepath, session_id=murfey_session.id, colors=colors, + denoising_suffix=denoising_suffix, ) if shuffle_message: preprocessing_messages.reverse() @@ -465,8 +468,8 @@ def test_run_with_db( ) assert len(ispyb_gs_search) == (len(preprocessing_messages) - 2) // 2 for gs in ispyb_gs_search: - # Check that all entries point to the denoised images ("_Lng_LVCC") - assert gs.gridSquareImage is not None and "_Lng_LVCC" in gs.gridSquareImage + # Check that all entries point to the denoised images + assert gs.gridSquareImage is not None and denoising_suffix in gs.gridSquareImage # Check that the GridSquare color flags and collection mode are set correctly for flag, value in color_flags.items():