Added logic to optionally insert collection mode and colour flags for the Atlas and GridSquare tables in ISPyB for the CLEM workflow

Author: tieneupin

src/murfey/server/ispyb.py

@@ -186,13 +186,23 @@ def do_update_atlas(
         atlas_image: str,
         pixel_size: float,
         slot: int | None,
+        collection_mode: str | None = None,
+        color_flags: dict[str, str | int] = {},
     ):
         try:
             with ISPyBSession() as db:
                 atlas = db.query(Atlas).filter(Atlas.atlasId == atlas_id).one()
                 atlas.atlasImage = atlas_image or atlas.atlasImage
                 atlas.pixelSize = pixel_size or atlas.pixelSize
                 atlas.cassetteSlot = slot or atlas.cassetteSlot
+                atlas.mode = collection_mode or atlas.mode
+                # Optionally insert colour flags if present
+                if color_flags:
+                    for (
+                        col_name,
+                        value,
+                    ) in color_flags.items():
+                        setattr(atlas, col_name, value)
                 db.add(atlas)
                 db.commit()
                 return {"success": True, "return_value": atlas.atlasId}
@@ -209,6 +219,7 @@ def do_insert_grid_square(
         atlas_id: int,
         grid_square_id: int,
         grid_square_parameters: GridSquareParameters,
+        color_flags: dict[str, int] = {},
     ):
         # most of this is for mypy
         if (
@@ -235,6 +246,10 @@ def do_insert_grid_square(
             stageLocationY=grid_square_parameters.y_stage_position,
             pixelSize=grid_square_parameters.pixel_size,
         )
+        # Optionally insert colour flags
+        if color_flags:
+            for col_name, value in color_flags.items():
+                setattr(record, col_name, value)
         try:
             with ISPyBSession() as db:
                 db.add(record)
@@ -250,7 +265,10 @@ def do_insert_grid_square(
         return {"success": False, "return_value": None}
 
     def do_update_grid_square(
-        self, grid_square_id: int, grid_square_parameters: GridSquareParameters
+        self,
+        grid_square_id: int,
+        grid_square_parameters: GridSquareParameters,
+        color_flags: dict[str, int] = {},
     ):
         try:
             with ISPyBSession() as db:
@@ -290,6 +308,12 @@ def do_update_grid_square(
                     grid_square.stageLocationY = grid_square_parameters.y_stage_position
                 if grid_square_parameters.pixel_size:
                     grid_square.pixelSize = grid_square_parameters.pixel_size
+                if grid_square_parameters.collection_mode:
+                    grid_square.mode = grid_square_parameters.collection_mode
+                # Optionally insert colour flags
+                if color_flags:
+                    for col_name, value in color_flags.items():
+                        setattr(grid_square, col_name, value)
                 db.add(grid_square)
                 db.commit()
                 return {"success": True, "return_value": grid_square.gridSquareId}

src/murfey/util/models.py

@@ -157,6 +157,9 @@ class GridSquareParameters(BaseModel):
     pixel_size: Optional[float] = None
     angle: Optional[float] = None
 
+    # Collection mode
+    collection_mode: Optional[str] = None
+
 
 class FoilHoleParameters(BaseModel):
     tag: str

src/murfey/workflows/clem/register_preprocessing_results.py

@@ -11,6 +11,7 @@
 import logging
 import re
 import traceback
+from collections.abc import Collection
 from importlib.metadata import entry_points
 from pathlib import Path
 from typing import Literal, Optional
@@ -186,6 +187,41 @@ def _register_clem_image_series(
     logger.info(f"CLEM preprocessing results registered for {result.series_name!r} ")
 
 
+color_columns = {
+    "gray": "hasGrey",
+    "red": "hasRed",
+    "green": "hasGreen",
+    "blue": "hasBlue",
+    "cyan": "hasCyan",
+    "magenta": "hasMagenta",
+    "yellow": "hasYellow",
+}
+
+
+def _get_color_flags(
+    colors: Collection[str] = [],
+):
+    return {
+        color_columns[color]: (1 if color in color_columns.keys() else 0)
+        for color in colors
+    }
+
+
+def _determine_collection_mode(
+    colors: Collection[str] = [],
+):
+    if not colors:
+        logger.warning("No colours were present in returned result")
+        return None
+    if "gray" in colors:
+        if len(colors) == 1:
+            return "Bright Field"
+        else:
+            return "Bright Field and Fluorescent"
+    else:
+        return "Fluorescent"
+
+
 def _register_dcg_and_atlas(
     session_id: int,
     instrument_name: str,
@@ -245,6 +281,10 @@ def _register_dcg_and_atlas(
                 "atlas": atlas_name,
                 "atlas_pixel_size": atlas_pixel_size,
                 "sample": dcg_entry.sample,
+                "color_flags": _get_color_flags(result.output_files.keys()),
+                "collection_mode": _determine_collection_mode(
+                    result.output_files.keys()
+                ),
             }
             if entry_point_result := entry_points(
                 group="murfey.workflows", name="atlas_update"
@@ -269,6 +309,8 @@ def _register_dcg_and_atlas(
             "atlas": atlas_name,
             "atlas_pixel_size": atlas_pixel_size,
             "sample": None,
+            "color_flags": _get_color_flags(result.output_files.keys()),
+            "collection_mode": _determine_collection_mode(result.output_files.keys()),
         }
         if entry_point_result := entry_points(
             group="murfey.workflows", name="data_collection_group"
@@ -410,6 +452,7 @@ def _register_grid_square(
                 y_stage_position=0.5 * (clem_img_series.y0 + clem_img_series.y1),
                 pixel_size=clem_img_series.image_pixel_size,
                 image=clem_img_series.thumbnail_search_string,
+                collection_mode=_determine_collection_mode(result.output_files.keys()),
             )
             # Register or update the grid square entry as required
             if grid_square_result := murfey_db.exec(

src/murfey/workflows/register_atlas_update.py

@@ -19,10 +19,13 @@ def run(
     logger.info(f"Registering updated atlas: \n{message}")
 
     _transport_object.do_update_atlas(
-        message["atlas_id"],
-        message["atlas"],
-        message["atlas_pixel_size"],
-        message["sample"],
+        atlas_id=message["atlas_id"],
+        atlas_image=message["atlas"],
+        pixel_size=message["atlas_pixel_size"],
+        slot=message["sample"],
+        # Extract optional parameters
+        collection_mode=message.get("collection_mode"),
+        color_flags=message.get("color_flags", {}),
     )
     if dcg_hooks := entry_points(group="murfey.hooks", name="data_collection_group"):
         try:

src/murfey/workflows/register_data_collection_group.py

@@ -67,6 +67,12 @@ def run(message: dict, murfey_db: SQLModelSession) -> dict[str, bool]:
                 pixelSize=message.get("atlas_pixel_size", 0),
                 cassetteSlot=message.get("sample"),
             )
+            # Optionally set the collection mode and color flags
+            if collection_mode := message.get("collection_mode"):
+                atlas_record.mode = collection_mode
+            if color_flags := message.get("color_flags", {}):
+                for col_name, value in color_flags.items():
+                    setattr(atlas_record, col_name, value)
             atlas_id = _transport_object.do_insert_atlas(atlas_record).get(
                 "return_value", None
             )