diff --git a/src/murfey/workflows/fib/register_lamella_evaluation_image.py b/src/murfey/workflows/fib/register_lamella_evaluation_image.py index 529fa8a06..ac3c7416d 100644 --- a/src/murfey/workflows/fib/register_lamella_evaluation_image.py +++ b/src/murfey/workflows/fib/register_lamella_evaluation_image.py @@ -16,6 +16,7 @@ parse_image_metadata, populate_fib_imaging_site_entry, ) +from murfey.workflows.register_data_collection_group import register_dcg logger = logging.getLogger(__name__) @@ -116,6 +117,67 @@ def _register_fib_imaging_site( return fib_imaging_site +def _register_dcg( + session_id: int, + instrument_name: str, + visit_name: str, + imaging_site: MurfeyDB.ImagingSite, + murfey_db: SQLModelSession, +): + """ + Takes an ImagingSite entry and uses it to create and register a DataCollectionGroup + entry in ISPyB if one doesn't already exist, or to populate an existing entry. + After doing so, it will register the DataCollectionGroup ID in Murfey and add it to + the ImagingSite entry. + """ + # Determine variables to register data collection group and atlas with + 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] + + # Generate a name/tag for the data collection group + # The name will be the site name minus the "/lamella..." specifier + dcg_name = "/".join(imaging_site.site_name.split("/")[:-1]) + + # Check if a DataCollectionGroup entry with this session and tag already exists + dcg_entry = murfey_db.exec( + select(MurfeyDB.DataCollectionGroup) + .where(MurfeyDB.DataCollectionGroup.session_id == session_id) + .where(MurfeyDB.DataCollectionGroup.tag == dcg_name) + ).one_or_none() + if not dcg_entry: + # Create a placeholder DataCollectionGroup and Atlas if not + dcg_message = { + "microscope": instrument_name, + "proposal_code": proposal_code, + "proposal_number": proposal_number, + "visit_number": visit_number, + "session_id": session_id, + "tag": dcg_name, + "experiment_type_id": 46, + "atlas": "", + "atlas_pixel_size": 0.0, + "sample": None, + } + dcg_entry = register_dcg( + message=dcg_message, + murfey_db=murfey_db, + ) + if not dcg_entry: + raise RuntimeError( + "Failed to create DataCollectionGroup entry for " + f"{imaging_site.image_path}" + ) + + # Update the ImagingSite with the DataCollectionGroup ID + imaging_site.dcg_id = dcg_entry.id + imaging_site.dcg_name = dcg_entry.tag + murfey_db.add(imaging_site) + murfey_db.commit() + + class FIBLamellaImageInfo(BaseModel): session_id: int lamella_image_file: Path @@ -129,84 +191,58 @@ def run( logger.info( f"Received the following message:\n{json.dumps(message, indent=2, default=str)}" ) - try: - try: - # Validate incoming message - fib_info = FIBLamellaImageInfo(**message) - except Exception: - logger.error("Could not validate incoming message", exc_info=True) - return {"success": False, "requeue": False} - - try: - # Load visit information - murfey_session = murfey_db.exec( - select(MurfeyDB.Session).where( - MurfeyDB.Session.id == fib_info.session_id - ) - ).one() - visit_name = murfey_session.visit - instrument_name = murfey_session.instrument_name - except Exception: - logger.error( - "Exception encountered while querying Murfey database", exc_info=True - ) - return {"success": False, "requeue": False} - try: - # Load the machine config - machine_config = get_machine_config(instrument_name)[instrument_name] - rotation_offset: float = cast( - float, machine_config.calibrations.get("rotation_offset", 0) - ) + # Validate incoming message + fib_info = FIBLamellaImageInfo(**message) - # Extract metadata from the image - metadata = FIBImageMetadata( - visit_name=visit_name, - file=fib_info.lamella_image_file, - **parse_image_metadata( - file=fib_info.lamella_image_file, - rotation_offset=rotation_offset, - ), - ) - logger.info( - "Extracted the following metadata from the image:\n" - f"{json.dumps(metadata.model_dump(), indent=2, default=str)}" - ) - except Exception: - logger.error( - f"Error extracting metadata from file {fib_info.lamella_image_file}", - exc_info=True, - ) - return {"success": False, "requeue": False} - - try: - # Make a thumbnail of the image and update metadata accordingly - metadata.thumbnail_path = _make_thumbnail( - file=fib_info.lamella_image_file, - metadata=metadata, - visit_name=visit_name, - ) - except Exception: - logger.warning( - f"Error creating thumbnail of file {fib_info.lamella_image_file}", - exc_info=True, - ) + # Load visit information + murfey_session = murfey_db.exec( + select(MurfeyDB.Session).where(MurfeyDB.Session.id == fib_info.session_id) + ).one() + visit_name = murfey_session.visit + instrument_name = murfey_session.instrument_name - try: - # Register imaging site to Murfey, or update existing one - _ = _register_fib_imaging_site(fib_info.session_id, metadata, murfey_db) - logger.info( - f"Registered lamella evaluation image {fib_info.lamella_image_file} " - f"for slot {metadata.slot_number} in Murfey database" - ) - except Exception: - logger.error( - "Error registering lamella evaluation image " - f"{fib_info.lamella_image_file} in Murfey database", - exc_info=True, - ) - return {"success": False, "requeue": False} + # Load the machine config + machine_config = get_machine_config(instrument_name)[instrument_name] + rotation_offset: float = cast( + float, machine_config.calibrations.get("rotation_offset", 0) + ) + + # Extract metadata from the image + metadata = FIBImageMetadata( + visit_name=visit_name, + file=fib_info.lamella_image_file, + **parse_image_metadata( + file=fib_info.lamella_image_file, + rotation_offset=rotation_offset, + ), + ) + logger.info( + "Extracted the following metadata from the image:\n" + f"{json.dumps(metadata.model_dump(), indent=2, default=str)}" + ) + + # Make a thumbnail of the image and update metadata accordingly + metadata.thumbnail_path = _make_thumbnail( + file=fib_info.lamella_image_file, + metadata=metadata, + visit_name=visit_name, + ) + + # Register imaging site to Murfey, or update existing one + fib_img_site = _register_fib_imaging_site(fib_info.session_id, metadata, murfey_db) + logger.info( + f"Registered lamella evaluation image {fib_info.lamella_image_file} " + f"for slot {metadata.slot_number} in Murfey database" + ) + + # Register data collection group and atlas in ISPyB + _register_dcg( + session_id=fib_info.session_id, + instrument_name=instrument_name, + visit_name=visit_name, + imaging_site=fib_img_site, + murfey_db=murfey_db, + ) - return {"success": True} - finally: - murfey_db.close() + return {"success": True} diff --git a/tests/workflows/fib/test_register_atlas.py b/tests/workflows/fib/test_register_atlas.py index 358d27f38..9f69074e8 100644 --- a/tests/workflows/fib/test_register_atlas.py +++ b/tests/workflows/fib/test_register_atlas.py @@ -11,7 +11,9 @@ from sqlmodel import Session as SQLModelSession, select as sm_select import murfey.util.db as MurfeyDB -from murfey.util.fib import get_slot_number +import murfey.workflows.fib.register_atlas +from murfey.server.ispyb import TransportManager +from murfey.util.fib import get_slot_number, number_from_name from murfey.util.models import FIBImageMetadata from murfey.workflows.fib.register_atlas import run from tests.conftest import ExampleVisit @@ -44,16 +46,16 @@ def test_run_with_db( # Add a test visit to the database if not ( - session_entry := murfey_db_session.exec( + murfey_session := murfey_db_session.exec( sm_select(MurfeyDB.Session).where(MurfeyDB.Session.id == session_id) ).one_or_none() ): - session_entry = MurfeyDB.Session(id=session_id) - session_entry.name = visit_name - session_entry.visit = visit_name - session_entry.instrument_name = instrument_name + murfey_session = MurfeyDB.Session(id=session_id) + murfey_session.name = visit_name + murfey_session.visit = visit_name + murfey_session.instrument_name = instrument_name - murfey_db_session.add(session_entry) + murfey_db_session.add(murfey_session) murfey_db_session.commit() # Mock the MachineConfig @@ -88,15 +90,11 @@ def test_run_with_db( ) # Patch the TransportManager object in the workflows called - from murfey.server.ispyb import TransportManager - mocker.patch( "murfey.server._transport_object", new=TransportManager("PikaTransport") ) # Mock the metadata returned from the image file - import murfey.workflows.fib.register_atlas - extracted = { "voltage": 2000, "shift_x": 0, @@ -109,17 +107,18 @@ def test_run_with_db( "rotation": -1.309, "tilt_alpha": 0.8, "tilt_beta": 0, - "pixels_x": 3072, - "pixels_y": 2048, + "pixels_x": 1500, + "pixels_y": 1000, "pixel_size_x": 1e-6, "pixel_size_y": 1e-6, } - extracted["slot_number"] = get_slot_number( + slot_number = get_slot_number( x=extracted["pos_x"], y=extracted["pos_y"], rotation=extracted["rotation"], rotation_offset=rotation_offset, ) + extracted["slot_number"] = slot_number mock_metadata = [ FIBImageMetadata( visit_name=visit_name, @@ -132,15 +131,36 @@ def test_run_with_db( "murfey.workflows.fib.register_atlas.parse_image_metadata", return_value=extracted, ) - spy_register = mocker.spy( - murfey.workflows.fib.register_atlas, - "_register_fib_imaging_site", - ) # Mock 'PIL.Image.open' and create a test image mock_open = mocker.patch("murfey.workflows.fib.register_atlas.PIL.Image.open") mock_open.__enter__.return_value = PIL.Image.fromarray( - np.ones((2048, 1152), dtype=np.uint8) + np.ones((1500, 1000), dtype=np.uint8) + ) + # Build the name of the expected test images + thumbnails: list[Path] = [] + for file in test_files: + image_number = number_from_name(file.stem) + thumbnail = ( + visit_dir + / "processed" + / visit_name + / f"grid_{slot_number}" + / "atlas" + / f"atlas_{str(image_number).zfill(2)}.png" + ) + thumbnail.parent.mkdir(parents=True, exist_ok=True) + thumbnail.touch(exist_ok=True) + thumbnails.append(thumbnail) + + # Set up spies for the different steps in 'run()' + spy_thumbnail = mocker.spy( + murfey.workflows.fib.register_atlas, + "_make_thumbnail", + ) + 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 @@ -154,6 +174,9 @@ def test_run_with_db( murfey_db=murfey_db_session, ) assert mock_parse.call_count == len(test_files) + assert spy_thumbnail.call_count == len(test_files) + for thumbnail in thumbnails: + assert thumbnail.is_file() assert spy_register.call_count == len(test_files) # Murfey's ImagingSite should have an entry diff --git a/tests/workflows/fib/test_register_lamella_evaluation_image.py b/tests/workflows/fib/test_register_lamella_evaluation_image.py index ecd4c8380..0d67c0fbe 100644 --- a/tests/workflows/fib/test_register_lamella_evaluation_image.py +++ b/tests/workflows/fib/test_register_lamella_evaluation_image.py @@ -1,12 +1,18 @@ from pathlib import Path +from unittest.mock import MagicMock import numpy as np import PIL.Image import pytest +from ispyb.sqlalchemy import _auto_db_schema as ISPyBDB from pytest_mock import MockerFixture -from sqlmodel import Session as SQLModelSession, 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 +import murfey.workflows.fib.register_lamella_evaluation_image +from murfey.server.ispyb import TransportManager from murfey.util.config import MachineConfig from murfey.workflows.fib.register_lamella_evaluation_image import ( FIBImageMetadata, @@ -127,7 +133,7 @@ def test_register_fib_imaging_site_with_db( ) # Only one entry should exist - found_sites = murfey_db_session.exec(select(MurfeyDB.ImagingSite)).all() + found_sites = murfey_db_session.exec(sm_select(MurfeyDB.ImagingSite)).all() assert len(found_sites) == 1 # Key parameters should be populated @@ -145,15 +151,20 @@ def test_run_with_db( mocker: MockerFixture, visit_dir: Path, murfey_db_session: SQLModelSession, + ispyb_db_session: SQLAlchemySession, + mock_ispyb_credentials, ): # Register a Session for this test - murfey_session = MurfeyDB.Session( - id=session_id, - visit=visit_name, - name=visit_name, - instrument_name=instrument_name, - started=True, - ) + if not ( + murfey_session := murfey_db_session.exec( + sm_select(MurfeyDB.Session).where(MurfeyDB.Session.id == session_id) + ).one_or_none() + ): + murfey_session = MurfeyDB.Session(id=session_id) + murfey_session.name = visit_name + murfey_session.visit = visit_name + murfey_session.instrument_name = instrument_name + murfey_db_session.add(murfey_session) murfey_db_session.commit() @@ -168,6 +179,27 @@ def test_run_with_db( return_value={instrument_name: machine_config}, ) + # Mock the ISPyB connection where the TransportManager class is located + mocker.patch( + "murfey.server.ispyb.get_security_config", + return_value=MagicMock(ispyb_credentials=mock_ispyb_credentials), + ) + 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 + mocker.patch( + "murfey.server._transport_object", new=TransportManager("PikaTransport") + ) + # Create the test image files and their thumbnails raw_lamella_dir = ( visit_dir @@ -219,7 +251,7 @@ def test_run_with_db( "pixel_size_x": 1e-6, "pixel_size_y": 1e-6, } - mocker.patch( + mock_parse = mocker.patch( "murfey.workflows.fib.register_lamella_evaluation_image.parse_image_metadata", return_value=metadata_dict, ) @@ -232,6 +264,16 @@ def test_run_with_db( np.ones((1500, 1000), dtype=np.uint8) ) + # Set up spies for the functions called by 'run()' + spy_thumbnail = mocker.spy( + murfey.workflows.fib.register_lamella_evaluation_image, + "_make_thumbnail", + ) + spy_register = mocker.spy( + murfey.workflows.fib.register_lamella_evaluation_image, + "_register_fib_imaging_site", + ) + # Run function and check that expected calls were made for file in files: # Construct the message to pass to the function @@ -244,7 +286,9 @@ def test_run_with_db( assert result["success"] # 'PIL.Image.open' should have been called for each image - assert mock_open.call_count == len(files) + assert mock_parse.call_count == len(files) + assert spy_thumbnail.call_count == len(files) + assert spy_register.call_count == len(files) # Both thumbnails should have been generated for thumbnail in thumbnails: @@ -252,7 +296,7 @@ def test_run_with_db( # There should only be one ImagingSite entry associated with the visit imaging_sites = murfey_db_session.exec( - select(MurfeyDB.ImagingSite) + sm_select(MurfeyDB.ImagingSite) .where(MurfeyDB.ImagingSite.session_id == session_id) .where(MurfeyDB.ImagingSite.data_type == "grid_square") ).all() @@ -270,3 +314,44 @@ def test_run_with_db( # Site name should have been constructed correctly assert imaging_site.site_name == f"{visit_name}/grid_2/lamella_1" + assert imaging_site.dcg_name == f"{visit_name}/grid_2" + + # 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 + + # Check that the Murfey DataCollectionGroup entry was populated correctly + murfey_dcg = murfey_dcg_search[0] + assert murfey_dcg.tag == f"{visit_name}/grid_2" + + # ISPyB's DataCollectionGroup should have an entry + 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 + + # Check that the ISPyB DataCollectionGroup entry was populated correctly + ispyb_dcg = ispyb_dcg_search[0] + assert ispyb_dcg.experimentTypeId == 46 + + # ISPyB's Atlas should have an entry + 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