Register FIB data in ISPyB

src/murfey/client/contexts/fib.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import logging
-import re
 import threading
 import xml.etree.ElementTree as ET
 from dataclasses import dataclass, field
@@ -11,6 +10,7 @@
 from murfey.client.context import Context
 from murfey.client.instance_environment import MurfeyInstanceEnvironment
 from murfey.util.client import capture_post
+from murfey.util.fib import number_from_name
 from murfey.util.models import (
     LamellaSiteInfo,
     MillingStepInfo,
@@ -24,22 +24,6 @@
 lock = threading.Lock()
 
 
-def _number_from_name(name: str) -> int:
-    """
-    In the AutoTEM and Maps workflows for the FIB, the sites and images are
-    auto-incremented with parenthesised numbers (e.g. "Lamella (2)"), with
-    the first site/image typically not having a number.
-
-    This function extracts the number from the file name, and returns 1 if
-    no such number is found.
-    """
-    return (
-        int(match.group(1))
-        if (match := re.search(r"^[\w\s]+\((\d+)\)$", name)) is not None
-        else 1
-    )
-
-
 T = TypeVar("T")
 
 
@@ -416,7 +400,7 @@ def _parse_autotem_metadata(self, file: Path):
             if (site_name := _parse_xml_text(site, "Name", str)) is None:
                 logger.warning("Current site doesn't have a name")
                 continue
-            site_num = _number_from_name(site_name)
+            site_num = number_from_name(site_name)
             site_info = LamellaSiteInfo(
                 project_name=project_name,
                 site_name=site_name,
@@ -555,7 +539,7 @@ def _make_drift_correction_gif(
         parts = file.parts
         try:
             lamella_name = parts[parts.index("Sites") + 1]
-            lamella_number = _number_from_name(lamella_name)
+            lamella_number = number_from_name(lamella_name)
         except Exception:
             logger.warning(
                 f"Could not extract metadata from file {file}", exc_info=True

src/murfey/util/db.py

@@ -10,7 +10,9 @@
 from sqlmodel import Enum, Field, Relationship, SQLModel, create_engine
 
 """
+=======================================================================================
 GENERAL
+=======================================================================================
 """
 
 
@@ -173,7 +175,9 @@ class ImagingSite(SQLModel, table=True):  # type: ignore
 
 
 """
+=======================================================================================
 TEM SESSION AND PROCESSING WORKFLOW
+=======================================================================================
 """
 
 
@@ -1072,7 +1076,9 @@ class CryoemInitialModel(SQLModel, table=True):  # type: ignore
 
 
 """
+=======================================================================================
 FUNCTIONS
+=======================================================================================
 """
 
 

src/murfey/util/fib.py

@@ -0,0 +1,21 @@
+"""
+General functinos specific to the FIB workflow
+"""
+
+import re
+
+
+def number_from_name(name: str) -> int:
+    """
+    In the AutoTEM and Maps workflows for the FIB, the sites and images are
+    auto-incremented with parenthesised numbers (e.g. "Lamella (2)"), with
+    the first site/image typically not having a number.
+
+    This function extracts the number from the file name, and returns 1 if
+    no such number is found.
+    """
+    return (
+        int(match.group(1))
+        if (match := re.search(r"^[\w\s]+\((\d+)\)$", name)) is not None
+        else 1
+    )

src/murfey/workflows/fib/register_atlas.py

@@ -1,6 +1,8 @@
 import logging
+import traceback
 import xml.etree.ElementTree as ET
 from functools import cached_property
+from importlib.metadata import entry_points
 from pathlib import Path
 
 import numpy as np
@@ -9,6 +11,7 @@
 from sqlmodel import Session, select
 
 import murfey.util.db as MurfeyDB
+from murfey.util.fib import number_from_name
 
 logger = logging.getLogger("murfey.workflows.fib.register_atlas")
 
@@ -150,36 +153,147 @@ def _register_fib_imaging_site(
     """
     Register FIB atlas in Murfey database or update existing entry.
     """
-    # Create new entry if one doesn't already exist
-    if not (
+
+    def _update_entry(
+        imaging_site: MurfeyDB.ImagingSite,
+        metadata: FIBAtlasMetadata,
+    ):
+        imaging_site.image_path = str(metadata.file)
+        imaging_site.pos_x = metadata.pos_x
+        imaging_site.pos_y = metadata.pos_y
+        imaging_site.pos_z = metadata.pos_z
+        imaging_site.rotation = float(np.rad2deg(metadata.rotation))
+        imaging_site.tilt_alpha = float(np.rad2deg(metadata.tilt_alpha))
+        imaging_site.tilt_beta = float(np.rad2deg(metadata.tilt_beta))
+        imaging_site.len_x = metadata.len_x
+        imaging_site.len_y = metadata.len_y
+        imaging_site.image_pixels_x = metadata.pixels_x
+        imaging_site.image_pixels_y = metadata.pixels_y
+        imaging_site.image_pixel_size = metadata.pixel_size
+
+        return imaging_site
+
+    if (
         fib_imaging_site := murfey_db.exec(
             select(MurfeyDB.ImagingSite)
             .where(MurfeyDB.ImagingSite.session_id == session_id)
-            .where(MurfeyDB.ImagingSite.image_path == str(metadata.file))
+            .where(MurfeyDB.ImagingSite.site_name == metadata.site_name)
+            .where(MurfeyDB.ImagingSite.data_type == "atlas")
         ).one_or_none()
-    ):
+    ) is None:
+        # Create new entry if one doesn't already exist
         fib_imaging_site = MurfeyDB.ImagingSite(
             session_id=session_id,
+            site_name=metadata.site_name,
             image_path=str(metadata.file),
             data_type="atlas",
         )
-    # Add/update entries
-    fib_imaging_site.site_name = metadata.site_name
-    fib_imaging_site.pos_x = metadata.pos_x
-    fib_imaging_site.pos_y = metadata.pos_y
-    fib_imaging_site.pos_z = metadata.pos_z
-    fib_imaging_site.rotation = float(np.rad2deg(metadata.rotation))
-    fib_imaging_site.tilt_alpha = float(np.rad2deg(metadata.tilt_alpha))
-    fib_imaging_site.tilt_beta = float(np.rad2deg(metadata.tilt_beta))
-    fib_imaging_site.len_x = metadata.len_x
-    fib_imaging_site.len_y = metadata.len_y
-    fib_imaging_site.image_pixels_x = metadata.pixels_x
-    fib_imaging_site.image_pixels_y = metadata.pixels_y
-    fib_imaging_site.image_pixel_size = metadata.pixel_size
+        fib_imaging_site = _update_entry(fib_imaging_site, metadata)
+    else:
+        # Check if the entry is new or newer than the current stored one
+        incoming_number = number_from_name(metadata.file.stem)
+        # Handle empty string
+        if not fib_imaging_site.image_path:
+            current_number = 0
+        # Read 'maps' atlases in one way
+        elif "maps" in (curr_path := Path(fib_imaging_site.image_path)).parts:
+            current_number = number_from_name(curr_path.stem)
+        else:
+            current_number = 0
+        # Update if incoming one is newer
+        if incoming_number > current_number:
+            fib_imaging_site = _update_entry(fib_imaging_site, metadata)
 
     murfey_db.add(fib_imaging_site)
     murfey_db.commit()
 
+    return fib_imaging_site
+
+
+def _register_dcg_and_atlas(
+    session_id: int,
+    instrument_name: str,
+    visit_name: str,
+    imaging_site: MurfeyDB.ImagingSite,
+    metadata: FIBAtlasMetadata,
+    murfey_db: Session,
+):
+    proposal_code = "".join(char for char in visit_name.split("-")[0] if char.isalpha())
+    proposal_number = "".join(
+        char for char in visit_name.split("-")[0] if char.isdigit()
+    )
+    visit_number = visit_name.split("-")[-1]
+
+    # Register using thumbnail values if they are provided
+    if (
+        imaging_site.thumbnail_path is not None
+        and imaging_site.thumbnail_pixel_size is not None
+    ):
+        atlas_name: str | None = imaging_site.thumbnail_path
+        atlas_pixel_size: float | None = imaging_site.thumbnail_pixel_size
+    else:
+        atlas_name = imaging_site.image_path
+        atlas_pixel_size = imaging_site.image_pixel_size
+
+    if dcg_search := murfey_db.exec(
+        select(MurfeyDB.DataCollectionGroup)
+        .where(MurfeyDB.DataCollectionGroup.session_id == session_id)
+        .where(MurfeyDB.DataCollectionGroup.tag == imaging_site.site_name)
+    ).all():
+        dcg_entry = dcg_search[0]
+        atlas_message = {
+            "session_id": session_id,
+            "dcgid": dcg_entry.id,
+            "atlas_id": dcg_entry.atlas_id,
+            "atlas": atlas_name,
+            "atlas_pixel_size": atlas_pixel_size,
+            "sample": dcg_entry.sample,
+        }
+        if entry_point_result := entry_points(
+            group="murfey.workflows", name="atlas_update"
+        ):
+            (workflow,) = entry_point_result
+            _ = workflow.load()(
+                message=atlas_message,
+                murfey_db=murfey_db,
+            )
+        else:
+            logger.warning("No workflow found for 'atlas_update'")
+    else:
+        dcg_message = {
+            "microscope": instrument_name,
+            "proposal_code": proposal_code,
+            "proposal_number": proposal_number,
+            "visit_number": visit_number,
+            "session_id": session_id,
+            "tag": imaging_site.site_name,
+            "experiment_type_id": 46,
+            "atlas": atlas_name,
+            "atlas_pixel_size": atlas_pixel_size,
+            "sample": metadata.slot_number,
+        }
+        if entry_point_result := entry_points(
+            group="murfey.workflows", name="data_collection_group"
+        ):
+            (workflow,) = entry_point_result
+            # Register grid square
+            _ = workflow.load()(
+                message=dcg_message,
+                murfey_db=murfey_db,
+            )
+        else:
+            logger.warning("No workflow found for 'data_collection_group'")
+    dcg_entry = murfey_db.exec(
+        select(MurfeyDB.DataCollectionGroup)
+        .where(MurfeyDB.DataCollectionGroup.session_id == session_id)
+        .where(MurfeyDB.DataCollectionGroup.tag == imaging_site.site_name)
+    ).one()
+
+    imaging_site.dcg_id = dcg_entry.id
+    imaging_site.dcg_name = dcg_entry.tag
+    murfey_db.add(imaging_site)
+    murfey_db.commit()
+
 
 def run(
     session_id: int,
@@ -188,28 +302,30 @@ def run(
 ):
     # Outer try-finally block to ensure database connection closes
     try:
-        # Load visit information
         try:
-            session_entry = murfey_db.exec(
+            # Load visit information
+            murfey_session = murfey_db.exec(
                 select(MurfeyDB.Session).where(MurfeyDB.Session.id == session_id)
             ).one()
-            visit_name = session_entry.visit
+            visit_name = murfey_session.visit
         except Exception:
             logger.error(
                 "Exception encountered while querying Murfey database", exc_info=True
             )
             return False
 
-        # Extract metadata from Electron Snapshot image
         try:
+            # Extract metadata from Electron Snapshot image
             metadata = _parse_metadata(file, visit_name)
         except Exception:
             logger.error(f"Error extracting metadata from file {file}", exc_info=True)
             return False
 
-        # Register imaging site in Murfey, or update existing one
         try:
-            _register_fib_imaging_site(session_id, metadata, murfey_db)
+            # Register imaging site in Murfey, or update existing one
+            fib_imaging_site = _register_fib_imaging_site(
+                session_id, metadata, murfey_db
+            )
             logger.info(
                 f"Registered FIB atlas image {file} for slot {metadata.slot_number} in Murfey database"
             )
@@ -219,6 +335,26 @@ def run(
                 exc_info=True,
             )
             return False
+
+        try:
+            # Register data collection group and atlas in ISPyB
+            _register_dcg_and_atlas(
+                session_id=session_id,
+                instrument_name=murfey_session.instrument_name,
+                visit_name=murfey_session.visit,
+                imaging_site=fib_imaging_site,
+                metadata=metadata,
+                murfey_db=murfey_db,
+            )
+        except Exception:
+            # Log error but allow workflow to proceed
+            logger.error(
+                "Exception encountered when registering data collection group for FIB workflow "
+                f"for {metadata.site_name!r}: \n"
+                f"{traceback.format_exc()}"
+            )
+
         return True
+
     finally:
         murfey_db.close()

tests/client/contexts/test_fib.py

@@ -14,7 +14,6 @@
     FIBImage,
     _file_transferred_to,
     _get_source,
-    _number_from_name,
     _parse_boolean,
 )
 from murfey.util.models import LamellaSiteInfo
@@ -392,24 +391,6 @@ def fib_maps_images(visit_dir: Path):
 # -------------------------------------------------------------------------------------
 
 
-@pytest.mark.parametrize(
-    "test_params",
-    (  # File name | Expected number
-        # AutoTEM examples
-        ("Lamella", 1),
-        ("Lamella (2)", 2),
-        ("Lamella (12)", 12),
-        # Maps examples
-        ("Electron Snapshot", 1),
-        ("Electron Snapshot (3)", 3),
-        ("Electron Snapshot (21)", 21),
-    ),
-)
-def test_number_from_name(test_params: tuple[str, int]):
-    name, number = test_params
-    assert _number_from_name(name) == number
-
-
 @pytest.mark.parametrize(
     "test_params",
     (  # Input | Expected output