diff --git a/src/murfey/client/contexts/fib.py b/src/murfey/client/contexts/fib.py index 874fc6b94..0b06370d2 100644 --- a/src/murfey/client/contexts/fib.py +++ b/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 diff --git a/src/murfey/util/db.py b/src/murfey/util/db.py index 5ba6691d6..eb92f48fa 100644 --- a/src/murfey/util/db.py +++ b/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 +======================================================================================= """ diff --git a/src/murfey/util/fib.py b/src/murfey/util/fib.py new file mode 100644 index 000000000..7a5efdadb --- /dev/null +++ b/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 + ) diff --git a/src/murfey/workflows/fib/register_atlas.py b/src/murfey/workflows/fib/register_atlas.py index ebed1a198..251b97616 100644 --- a/src/murfey/workflows/fib/register_atlas.py +++ b/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() diff --git a/tests/client/contexts/test_fib.py b/tests/client/contexts/test_fib.py index f13c552f4..218a669c5 100644 --- a/tests/client/contexts/test_fib.py +++ b/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 diff --git a/tests/util/test_fib_util.py b/tests/util/test_fib_util.py new file mode 100644 index 000000000..6a8d9a73b --- /dev/null +++ b/tests/util/test_fib_util.py @@ -0,0 +1,21 @@ +import pytest + +from murfey.util.fib import number_from_name + + +@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 diff --git a/tests/workflows/fib/test_register_atlas.py b/tests/workflows/fib/test_register_atlas.py index 3b2600b51..27f0af52a 100644 --- a/tests/workflows/fib/test_register_atlas.py +++ b/tests/workflows/fib/test_register_atlas.py @@ -2,16 +2,20 @@ from pathlib import Path from unittest.mock import MagicMock +import ispyb.sqlalchemy as ISPyBDB import pytest from pytest_mock import MockerFixture -from sqlmodel import Session, select +from sqlalchemy import select as sa_select +from sqlalchemy.orm import Session as SQLAlchemySession +from sqlmodel import Session as SQLModelSession, select as sm_select import murfey.util.db as MurfeyDB from murfey.workflows.fib.register_atlas import FIBAtlasMetadata, _parse_metadata, run +from tests.conftest import ExampleVisit session_id = 10 -visit_name = "cm12345-6" -instrument_name = "test_instrument" +visit_name = f"{ExampleVisit.proposal_code}{ExampleVisit.proposal_number}-{ExampleVisit.visit_number}" +instrument_name = ExampleVisit.instrument_name @pytest.fixture @@ -289,16 +293,20 @@ def test_register_fib_imaging_site(): def test_run_with_db( mocker: MockerFixture, visit_dir: Path, - murfey_db_session: Session, + murfey_db_session: SQLModelSession, + ispyb_db_session: SQLAlchemySession, + mock_ispyb_credentials, ): - test_file = ( - visit_dir / "maps/LayersData/Layer/Electron Snapshot/Electron Snapshot.tiff" + test_files = ( + visit_dir / "maps/LayersData/Layer/Electron Snapshot/Electron Snapshot.tiff", + visit_dir + / "maps/LayersData/Layer/Electron Snapshot/Electron Snapshot (2).tiff", ) # Add a test visit to the database if not ( session_entry := murfey_db_session.exec( - select(MurfeyDB.Session).where(MurfeyDB.Session.id == session_id) + sm_select(MurfeyDB.Session).where(MurfeyDB.Session.id == session_id) ).one_or_none() ): session_entry = MurfeyDB.Session(id=session_id) @@ -309,34 +317,119 @@ def test_run_with_db( murfey_db_session.add(session_entry) murfey_db_session.commit() - # Mock the metadata returned from the image file - mock_metadata = FIBAtlasMetadata( - visit_name=visit_name, - file=test_file, - voltage=2000, - shift_x=0, - shift_y=0, - len_x=0.003072, - len_y=0.002048, - pos_x=0.003, - pos_y=0.0003, - pos_z=0.01, - rotation=-1.309, - tilt_alpha=0.8, - tilt_beta=0, - pixels_x=3072, - pixels_y=2048, - pixel_size_x=1e-6, - pixel_size_y=1e-6, + # Mock the ISPyB connection where the TransportManager class is located + mock_security_config = MagicMock() + mock_security_config.ispyb_credentials = mock_ispyb_credentials + mocker.patch( + "murfey.server.ispyb.get_security_config", + return_value=mock_security_config, + ) + mocker.patch( + "murfey.server.ispyb.ISPyBSession", + return_value=ispyb_db_session, + ) + + # Mock the ISPYB connection when registering data collection group + mocker.patch( + "murfey.workflows.register_data_collection_group.ISPyBSession", + return_value=ispyb_db_session, + ) + + # Patch the TransportManager object in the workflows called + from murfey.server.ispyb import TransportManager + + mocker.patch( + "murfey.workflows.register_data_collection_group._transport_object", + new=TransportManager("PikaTransport"), ) mocker.patch( + "murfey.workflows.register_atlas_update._transport_object", + new=TransportManager("PikaTransport"), + ) + + # Mock the metadata returned from the image file + import murfey.workflows.fib.register_atlas + + mock_metadata = [ + FIBAtlasMetadata( + visit_name=visit_name, + file=test_file, + voltage=2000, + shift_x=0, + shift_y=0, + len_x=0.003072, + len_y=0.002048, + pos_x=0.003, + pos_y=0.0003, + pos_z=0.01, + rotation=-1.309, + tilt_alpha=0.8, + tilt_beta=0, + pixels_x=3072, + pixels_y=2048, + pixel_size_x=1e-6, + pixel_size_y=1e-6, + ) + for test_file in test_files + ] + mock_parse = mocker.patch( "murfey.workflows.fib.register_atlas._parse_metadata", - return_value=mock_metadata, + side_effect=mock_metadata, + ) + spy_register = mocker.spy( + murfey.workflows.fib.register_atlas, + "_register_fib_imaging_site", ) # Run the function and check that it's run through to completion - assert run( - session_id=session_id, - file=test_file, - murfey_db=murfey_db_session, + for test_file in test_files: + run( + session_id=session_id, + file=test_file, + murfey_db=murfey_db_session, + ) + assert mock_parse.call_count == len(test_files) + assert spy_register.call_count == len(test_files) + + # Murfey's ImagingSite should have an entry + search_results = murfey_db_session.exec( + sm_select(MurfeyDB.ImagingSite).where( + MurfeyDB.ImagingSite.session_id == session_id + ) + ).all() + assert len(search_results) == 1 + assert search_results[0].image_path == str(mock_metadata[-1].file) + + # Murfey's DataCollectionGroup should have an entry + murfey_dcg_search = murfey_db_session.exec( + sm_select(MurfeyDB.DataCollectionGroup).where( + MurfeyDB.DataCollectionGroup.session_id == session_id + ) + ).all() + assert len(murfey_dcg_search) == 1 + + # ISPyB's DataCollectionGroup should have an entry + murfey_dcg = murfey_dcg_search[0] + ispyb_dcg_search = ( + ispyb_db_session.execute( + sa_select(ISPyBDB.DataCollectionGroup).where( + ISPyBDB.DataCollectionGroup.dataCollectionGroupId == murfey_dcg.id + ) + ) + .scalars() + .all() + ) + assert len(ispyb_dcg_search) == 1 + + # Atlas should have an entry + ispyb_dcg = ispyb_dcg_search[0] + ispyb_atlas_search = ( + ispyb_db_session.execute( + sa_select(ISPyBDB.Atlas).where( + ISPyBDB.Atlas.dataCollectionGroupId == ispyb_dcg.dataCollectionGroupId + ) + ) + .scalars() + .all() ) + assert len(ispyb_atlas_search) == 1