diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 4f34779c4..bcf26da33 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -3,10 +3,11 @@ name: Build and test on: [push, pull_request] env: - ISPYB_DATABASE_SCHEMA: 4.11.0 + ISPYB_DATABASE_SCHEMA: 5.1.0 # Installs from GitHub # Versions: https://github.com/DiamondLightSource/ispyb-database/tags # Previous version(s): + # 4.11.0 # 4.8.0 # 4.2.1 # released 2024-08-19 # 4.1.0 # released 2024-03-26 diff --git a/pyproject.toml b/pyproject.toml index b181e111b..5e8e46951 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -55,7 +55,7 @@ developer = [ server = [ "cryptography", "graypy", - "ispyb>=10.2.4", # Responsible for setting requirements for SQLAlchemy and mysql-connector-python; + "ispyb>=12.1.0", # Responsible for setting requirements for SQLAlchemy and mysql-connector-python; "jinja2", "mrcfile", "numpy<3", @@ -117,6 +117,7 @@ TomographyMetadataContext = "murfey.client.contexts.tomo_metadata:TomographyMeta "data_collection_group" = "murfey.workflows.register_data_collection_group:run" "experiment_type_update" = "murfey.workflows.register_experiment_type_update:run" "fib.register_atlas" = "murfey.workflows.fib.register_atlas:run" +"fib.register_milling_progress" = "murfey.workflows.fib.register_milling_progress:run" "pato" = "murfey.workflows.notifications:notification_setup" "picked_particles" = "murfey.workflows.spa.picking:particles_picked" "picked_tomogram" = "murfey.workflows.tomo.picking:picked_tomogram" diff --git a/src/murfey/server/api/workflow_fib.py b/src/murfey/server/api/workflow_fib.py index 577a11a4d..5c143a1f5 100644 --- a/src/murfey/server/api/workflow_fib.py +++ b/src/murfey/server/api/workflow_fib.py @@ -1,4 +1,3 @@ -import json import logging import os from pathlib import Path @@ -7,7 +6,7 @@ import PIL.Image from fastapi import APIRouter, Depends from pydantic import BaseModel -from sqlmodel import Session, select +from sqlmodel import select import murfey.util.db as MurfeyDB from murfey.server import _transport_object @@ -36,7 +35,7 @@ def register_fib_atlas( fib_atlas: FIBAtlasFile, ): if _transport_object is None: - logger.error("No Transport Manager object was set up") + logger.error("No TransportManager object was set up") return None _transport_object.send( _transport_object.feedback_queue, @@ -52,11 +51,17 @@ def register_fib_atlas( def register_fib_milling_progress( session_id: int, site_info: LamellaSiteInfo, - db: Session = murfey_db, ): - logger.debug( - "Received the following FIB metadata for registration:\n" - f"{json.dumps(site_info.model_dump(exclude_none=True), indent=2, default=str)}" + if _transport_object is None: + logger.error("No TransportManager object was set up") + return None + _transport_object.send( + _transport_object.feedback_queue, + { + "register": "fib.register_milling_progress", + "session_id": session_id, + "site_info": site_info.model_dump(exclude_none=True), + }, ) diff --git a/src/murfey/server/ispyb.py b/src/murfey/server/ispyb.py index 815f06354..75525dd0c 100644 --- a/src/murfey/server/ispyb.py +++ b/src/murfey/server/ispyb.py @@ -20,6 +20,8 @@ DataCollectionGroup, FoilHole, GridSquare, + MillingStep, + MillingStepName, ProcessingJob, ProcessingJobParameter, Proposal, @@ -710,6 +712,164 @@ def do_update_processing_status(self, record: AutoProcProgram, **kwargs): ) return {"success": False, "return_value": None} + def do_insert_milling_step( + self, + # MillingStepName identifiers + recipe_name: str, + activity_name: str, + grid_square_id: int, + # Values + is_enabled: bool | None = None, + status: str | None = None, + execution_time: float | None = None, + stage_x: float | None = None, + stage_y: float | None = None, + stage_z: float | None = None, + rotation: float | None = None, + tilt_alpha: float | None = None, + beam_type: str | None = None, + beam_voltage: float | None = None, + beam_current: float | None = None, + milling_angle: float | None = None, + depth_correction: float | None = None, + lamella_offset: float | None = None, + trench_height_front: float | None = None, + trench_height_rear: float | None = None, + width_overlap_front_left: float | None = None, + width_overlap_front_right: float | None = None, + width_overlap_rear_left: float | None = None, + width_overlap_rear_right: float | None = None, + ): + try: + with ISPyBSession() as db: + # Find the ID of this MillingStep + milling_step_name = ( + db.query(MillingStepName) + .filter( + MillingStepName.recipe == recipe_name, + MillingStepName.step == activity_name, + ) + .one() + ) + record = MillingStep( + # IDs + millingStepNameId=milling_step_name.millingStepNameId, + gridSquareId=grid_square_id, + # Values + isEnabled=is_enabled, + status=status, + executionTime=execution_time, + stageX=stage_x, + stageY=stage_y, + stageZ=stage_z, + rotation=rotation, + alphaTilt=tilt_alpha, + beamType=beam_type, + beamVoltage=beam_voltage, + beamCurrent=beam_current, + millingAngle=milling_angle, + depthCorrection=depth_correction, + lamellaOffset=lamella_offset, + trenchHeightFront=trench_height_front, + trenchHeightRear=trench_height_rear, + widthOverlapFrontLeft=width_overlap_front_left, + widthOverlapFrontRight=width_overlap_front_right, + widthOverlapRearLeft=width_overlap_rear_left, + widthOverlapRearRight=width_overlap_rear_right, + ) + db.add(record) + db.commit() + log.info(f"Created MillingStep {record.millingStepId}") + return {"success": True, "return_value": record.millingStepId} + except ispyb.ISPyBException as e: + log.error( + "Insert MillingStep entry caused exception '%s'.", + e, + exc_info=True, + ) + return {"success": False, "return_value": None} + + def do_update_milling_step( + self, + milling_step_id: int, + is_enabled: bool | None = None, + status: str | None = None, + execution_time: float | None = None, + stage_x: float | None = None, + stage_y: float | None = None, + stage_z: float | None = None, + rotation: float | None = None, + tilt_alpha: float | None = None, + beam_type: str | None = None, + beam_voltage: float | None = None, + beam_current: float | None = None, + milling_angle: float | None = None, + depth_correction: float | None = None, + lamella_offset: float | None = None, + trench_height_front: float | None = None, + trench_height_rear: float | None = None, + width_overlap_front_left: float | None = None, + width_overlap_front_right: float | None = None, + width_overlap_rear_left: float | None = None, + width_overlap_rear_right: float | None = None, + ): + try: + with ISPyBSession() as db: + milling_step = ( + db.query(MillingStep) + .filter(MillingStep.millingStepId == milling_step_id) + .one() + ) + milling_step.isEnabled = is_enabled or milling_step.isEnabled + milling_step.status = status or milling_step.status + milling_step.executionTime = ( + execution_time or milling_step.executionTime + ) + milling_step.stageX = stage_x or milling_step.stageX + milling_step.stageY = stage_y or milling_step.stageY + milling_step.stageZ = stage_z or milling_step.stageZ + milling_step.rotation = rotation or milling_step.rotation + milling_step.alphaTilt = tilt_alpha or milling_step.alphaTilt + milling_step.beamType = beam_type or milling_step.beamType + milling_step.beamVoltage = beam_voltage or milling_step.beamVoltage + milling_step.beamCurrent = beam_current or milling_step.beamCurrent + milling_step.millingAngle = milling_angle or milling_step.millingAngle + milling_step.depthCorrection = ( + depth_correction or milling_step.depthCorrection + ) + milling_step.lamellaOffset = ( + lamella_offset or milling_step.lamellaOffset + ) + milling_step.trenchHeightFront = ( + trench_height_front or milling_step.trenchHeightFront + ) + milling_step.trenchHeightRear = ( + trench_height_rear or milling_step.trenchHeightRear + ) + milling_step.widthOverlapFrontLeft = ( + width_overlap_front_left or milling_step.widthOverlapFrontLeft + ) + milling_step.widthOverlapFrontRight = ( + width_overlap_front_right or milling_step.widthOverlapFrontRight + ) + milling_step.widthOverlapRearLeft = ( + width_overlap_rear_left or milling_step.widthOverlapRearLeft + ) + milling_step.widthOverlapRearRight = ( + width_overlap_rear_right or milling_step.widthOverlapRearRight + ) + + db.add(milling_step) + db.commit() + return {"success": True, "return_value": milling_step.millingStepId} + except ispyb.ISPyBException as e: + log.error( + "Updating MillingStep entry caused exception '%s'.", + e, + exc_info=True, + ) + return {"success": False, "return_value": None} + def do_buffer_lookup(self, app_id: int, uuid: int) -> Optional[int]: with ISPyBSession() as db: buffer_objects = ( diff --git a/src/murfey/util/db.py b/src/murfey/util/db.py index eb92f48fa..413336694 100644 --- a/src/murfey/util/db.py +++ b/src/murfey/util/db.py @@ -174,6 +174,139 @@ class ImagingSite(SQLModel, table=True): # type: ignore grid_square_id: Optional[int] = Field(foreign_key="gridsquare.id", default=None) +class DataCollectionGroup(SQLModel, table=True): # type: ignore + id: int = Field( + primary_key=True, + unique=True, + alias="dataCollectionGroupId", + sa_column_kwargs={"name": "dataCollectionGroupId"}, + ) + session_id: int = Field(foreign_key="session.id", primary_key=True) + tag: str = Field(primary_key=True) + atlas_id: Optional[int] = None + atlas_pixel_size: Optional[float] = None + atlas: str = "" + sample: Optional[int] = None + smartem_grid_uuid: Optional[str] = None + + # ------------- + # Relationships + # ------------- + session: Optional["Session"] = Relationship(back_populates="data_collection_groups") + data_collections: List["DataCollection"] = Relationship( + back_populates="data_collection_group", + sa_relationship_kwargs={"cascade": "delete"}, + ) + imaging_sites: List["ImagingSite"] = Relationship( + back_populates="data_collection_group", + sa_relationship_kwargs={"cascade": "delete"}, + ) + notification_parameters: List["NotificationParameter"] = Relationship( + back_populates="data_collection_group", + sa_relationship_kwargs={"cascade": "delete"}, + ) + tomography_processing_parameters: List["TomographyProcessingParameters"] = ( + Relationship( + back_populates="data_collection_group", + sa_relationship_kwargs={"cascade": "delete"}, + ) + ) + grid_squares: Optional[List["GridSquare"]] = Relationship( + back_populates="data_collection_group", + sa_relationship_kwargs={"cascade": "delete"}, + ) + search_maps: Optional[List["SearchMap"]] = Relationship( + back_populates="data_collection_group", + sa_relationship_kwargs={"cascade": "delete"}, + ) + + +class GridSquare(SQLModel, table=True): # type: ignore + id: Optional[int] = Field(primary_key=True, default=None) + session_id: int = Field(foreign_key="session.id") + name: int + tag: str + x_location: Optional[float] + y_location: Optional[float] + x_stage_position: Optional[float] + y_stage_position: Optional[float] + readout_area_x: Optional[int] + readout_area_y: Optional[int] + thumbnail_size_x: Optional[int] + thumbnail_size_y: Optional[int] + pixel_size: Optional[float] = None + image: str = "" + atlas_id: Optional[int] = Field( + foreign_key="datacollectiongroup.dataCollectionGroupId" + ) + scaled_pixel_size: Optional[float] = None + pixel_location_x: Optional[int] = None + pixel_location_y: Optional[int] = None + height: Optional[int] = None + width: Optional[int] = None + angle: Optional[float] = None + quality_indicator: Optional[float] = None + smartem_uuid: Optional[str] = None + + # ------------- + # Relationships + # ------------- + session: Optional[Session] = Relationship(back_populates="grid_squares") + imaging_sites: List["ImagingSite"] = Relationship( + back_populates="grid_square", sa_relationship_kwargs={"cascade": "delete"} + ) + milling_steps: List["MillingStep"] = Relationship( + back_populates="grid_square", sa_relationship_kwargs={"cascade": "delete"} + ) + foil_holes: List["FoilHole"] = Relationship( + back_populates="grid_square", sa_relationship_kwargs={"cascade": "delete"} + ) + data_collection_group: Optional["DataCollectionGroup"] = Relationship( + back_populates="grid_squares" + ) + + +""" +======================================================================================= +FIB WORKFLOW +======================================================================================= +""" + + +class MillingStep(SQLModel, table=True): # type: ignore + id: int = Field(primary_key=True, default=None) + recipe_name: Optional[str] = Field(default=None) + activity_name: Optional[str] = Field(default=None) + is_enabled: Optional[bool] = Field(default=None) + status: Optional[str] = Field(default=None) + execution_time: Optional[float] = Field(default=None) + stage_x: Optional[float] = Field(default=None) + stage_y: Optional[float] = Field(default=None) + stage_z: Optional[float] = Field(default=None) + rotation: Optional[float] = Field(default=None) + tilt_alpha: Optional[float] = Field(default=None) + beam_type: Optional[str] = Field(default=None) + beam_voltage: Optional[float] = Field(default=None) + beam_current: Optional[float] = Field(default=None) + milling_angle: Optional[float] = Field(default=None) + depth_correction: Optional[float] = Field(default=None) + lamella_offset: Optional[float] = Field(default=None) + trench_height_front: Optional[float] = Field(default=None) + trench_height_rear: Optional[float] = Field(default=None) + width_overlap_front_left: Optional[float] = Field(default=None) + width_overlap_front_right: Optional[float] = Field(default=None) + width_overlap_rear_left: Optional[float] = Field(default=None) + width_overlap_rear_right: Optional[float] = Field(default=None) + + # ------------- + # Relationships + # ------------- + grid_square: Optional["GridSquare"] = Relationship( + back_populates="milling_steps" + ) # Many-to-one + grid_square_id: Optional[int] = Field(foreign_key="gridsquare.id", default=None) + + """ ======================================================================================= TEM SESSION AND PROCESSING WORKFLOW @@ -223,49 +356,6 @@ class Tilt(SQLModel, table=True): # type: ignore tilt_series: Optional[TiltSeries] = Relationship(back_populates="tilts") -class DataCollectionGroup(SQLModel, table=True): # type: ignore - id: int = Field( - primary_key=True, - unique=True, - alias="dataCollectionGroupId", - sa_column_kwargs={"name": "dataCollectionGroupId"}, - ) - session_id: int = Field(foreign_key="session.id", primary_key=True) - tag: str = Field(primary_key=True) - atlas_id: Optional[int] = None - atlas_pixel_size: Optional[float] = None - atlas: str = "" - sample: Optional[int] = None - smartem_grid_uuid: Optional[str] = None - session: Optional["Session"] = Relationship(back_populates="data_collection_groups") - data_collections: List["DataCollection"] = Relationship( - back_populates="data_collection_group", - sa_relationship_kwargs={"cascade": "delete"}, - ) - imaging_sites: List["ImagingSite"] = Relationship( - back_populates="data_collection_group", - sa_relationship_kwargs={"cascade": "delete"}, - ) - notification_parameters: List["NotificationParameter"] = Relationship( - back_populates="data_collection_group", - sa_relationship_kwargs={"cascade": "delete"}, - ) - tomography_processing_parameters: List["TomographyProcessingParameters"] = ( - Relationship( - back_populates="data_collection_group", - sa_relationship_kwargs={"cascade": "delete"}, - ) - ) - grid_squares: Optional[List["GridSquare"]] = Relationship( - back_populates="data_collection_group", - sa_relationship_kwargs={"cascade": "delete"}, - ) - search_maps: Optional[List["SearchMap"]] = Relationship( - back_populates="data_collection_group", - sa_relationship_kwargs={"cascade": "delete"}, - ) - - class NotificationParameter(SQLModel, table=True): # type: ignore id: Optional[int] = Field(default=None, primary_key=True) dcg_id: int = Field(foreign_key="datacollectiongroup.dataCollectionGroupId") @@ -482,44 +572,6 @@ class MurfeyLedger(SQLModel, table=True): # type: ignore ) -class GridSquare(SQLModel, table=True): # type: ignore - id: Optional[int] = Field(primary_key=True, default=None) - session_id: int = Field(foreign_key="session.id") - name: int - tag: str - x_location: Optional[float] - y_location: Optional[float] - x_stage_position: Optional[float] - y_stage_position: Optional[float] - readout_area_x: Optional[int] - readout_area_y: Optional[int] - thumbnail_size_x: Optional[int] - thumbnail_size_y: Optional[int] - pixel_size: Optional[float] = None - image: str = "" - session: Optional[Session] = Relationship(back_populates="grid_squares") - imaging_sites: List["ImagingSite"] = Relationship( - back_populates="grid_square", sa_relationship_kwargs={"cascade": "delete"} - ) - foil_holes: List["FoilHole"] = Relationship( - back_populates="grid_square", sa_relationship_kwargs={"cascade": "delete"} - ) - atlas_id: Optional[int] = Field( - foreign_key="datacollectiongroup.dataCollectionGroupId" - ) - scaled_pixel_size: Optional[float] = None - pixel_location_x: Optional[int] = None - pixel_location_y: Optional[int] = None - height: Optional[int] = None - width: Optional[int] = None - angle: Optional[float] = None - quality_indicator: Optional[float] = None - smartem_uuid: Optional[str] = None - data_collection_group: Optional["DataCollectionGroup"] = Relationship( - back_populates="grid_squares" - ) - - class FoilHole(SQLModel, table=True): # type: ignore id: Optional[int] = Field(primary_key=True, default=None) grid_square_id: int = Field(foreign_key="gridsquare.id") diff --git a/src/murfey/util/models.py b/src/murfey/util/models.py index f2f00fc0e..a3e594118 100644 --- a/src/murfey/util/models.py +++ b/src/murfey/util/models.py @@ -179,6 +179,13 @@ class MillingStepInfo(BaseModel): width_overlap_rear_left: float | None = None width_overlap_rear_right: float | None = None + @field_validator("status", mode="before") + @classmethod + def handle_stringified_none(cls, v: Any) -> None: + if isinstance(v, str) and v.lower() == "none": + return None + return v + class MillingSteps(BaseModel): # Processing steps supported by AutoTEM diff --git a/src/murfey/workflows/fib/register_milling_progress.py b/src/murfey/workflows/fib/register_milling_progress.py new file mode 100644 index 000000000..3c209dc38 --- /dev/null +++ b/src/murfey/workflows/fib/register_milling_progress.py @@ -0,0 +1,454 @@ +from __future__ import annotations + +import json +import logging +from importlib.metadata import entry_points +from typing import TYPE_CHECKING, Any + +from sqlmodel import Session as SQLModelSession, select + +import murfey.util.db as MurfeyDB +from murfey.server import _transport_object +from murfey.util.models import ( + GridSquareParameters, + LamellaSiteInfo, + MillingStepInfo, + MillingSteps, + StagePositionInfo, + StagePositionValues, +) + +if TYPE_CHECKING: + from murfey.server.ispyb import TransportManager + + +logger = logging.getLogger("murfey.workflows.fib.register_milling_progress") + + +def _ensure_prerequisites( + session_id: int, + instrument_name: str, + visit_name: str, + project_name: str, + slot_number: int, + site_number: int, + transport_object: TransportManager, + murfey_db: SQLModelSession, +): + """ + Uses the FIB milling metadata provided to create the necessary DataCollectionGroup, + Atlas, and GridSquare placeholders in Murfey if they don't already exist + + """ + # Construct the DataCollectionGroup and GridSquare lookup tags + dcg_tag = f"{project_name}--slot_{slot_number}" + + # 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] + + # Register the DataCollectionGroup and Atlas placeholder if it doesn't already exist + if ( + murfey_db.exec( + select(MurfeyDB.DataCollectionGroup) + .where(MurfeyDB.DataCollectionGroup.session_id == session_id) + .where(MurfeyDB.DataCollectionGroup.tag == dcg_tag) + ).one_or_none() + is None + ): + dcg_message = { + "microscope": instrument_name, + "proposal_code": proposal_code, + "proposal_number": proposal_number, + "visit_number": visit_number, + "session_id": session_id, + "tag": dcg_tag, + "experiment_type_id": 46, + "atlas": "", + "atlas_pixel_size": 0.0, + "sample": slot_number, + "color_flags": None, + "collection_mode": None, + } + if entry_point_result := entry_points( + group="murfey.workflows", name="data_collection_group" + ): + (workflow,) = entry_point_result + _ = workflow.load()( + message=dcg_message, + murfey_db=murfey_db, + ) + + # Register the GridSquare if it doesn't already exist + grid_square_entry = murfey_db.exec( + select(MurfeyDB.GridSquare) + .where(MurfeyDB.GridSquare.session_id == session_id) + .where(MurfeyDB.GridSquare.name == site_number) + .where(MurfeyDB.GridSquare.tag == dcg_tag) + ).one_or_none() + if grid_square_entry is None: + dcg_entry = murfey_db.exec( + select(MurfeyDB.DataCollectionGroup) + .where(MurfeyDB.DataCollectionGroup.session_id == session_id) + .where(MurfeyDB.DataCollectionGroup.tag == dcg_tag) + ).one() + grid_square_ispyb_result = transport_object.do_insert_grid_square( + atlas_id=dcg_entry.atlas_id, + grid_square_id=site_number, + grid_square_parameters=GridSquareParameters( + tag=dcg_tag, + ), + ) + # Register to Murfey + grid_square_entry = MurfeyDB.GridSquare( + id=grid_square_ispyb_result.get("return_value", None), + name=site_number, + session_id=session_id, + tag=dcg_tag, + ) + murfey_db.add(grid_square_entry) + murfey_db.commit() + + return grid_square_entry + + +MILLING_STEP_LOOKUP = ( + # Match the MillingStep key names to their ISPyB IDs + # Preparation stage + ( + ( + "eucentric_tilt", + "artificial_features", + "milling_angle", + "image_acquisition", + "lamella_placement", + ), + "preparation_site", + ), + # Milling stage + ( + ( + "delay_1", + "reference_definition", + "reference_definition_electron", + "stress_relief_cuts", + "reference_redefinition_1", + "rough_milling", + "rough_milling_electron", + "reference_redefinition_2", + "medium_milling", + "medium_milling_electron", + "fine_milling", + "fine_milling_electron", + "finer_milling", + "finer_milling_electron", + ), + "chunk_site", + ), + # Thinning stage + ( + ( + "delay_2", + "polishing_1", + "polishing_1_electron", + "polishing_2", + "polishing_2_ion", + "polishing_2_electron", + ), + "thinning_site", + ), +) + + +def _register_milling_step( + milling_steps: MillingSteps, + stage_info: StagePositionInfo, + grid_square: MurfeyDB.GridSquare, + transport_object: TransportManager, + murfey_db: SQLModelSession, +): + """ + Registers FIB milling metadata for the current lamella site as MillingStep entries + in ISPyB. If successful, will proceed to create a backup copy of the inserted row + in Murfey. + """ + # Check that GridSquare has ID (for type checking) + if grid_square.id is None: + logger.error("Current GridSquare entry has no ID") + return None + + # Iteratively go through the LamellaSiteInfo model and insert for each step + for steps, stage_name in MILLING_STEP_LOOKUP: + for step_name in steps: + step_info: MillingStepInfo | None = milling_steps.__getattribute__( + step_name + ) + # Early continues if key information is missing + if step_info is None: + logger.debug(f"No step info found for {step_name}") + continue + if step_info.recipe_name is None: + logger.debug(f"No recipe name found for {step_name}") + continue + if step_info.step_name is None: + logger.debug(f"No step name found for {step_name}") + continue + + stage_values: StagePositionValues | None = stage_info.__getattribute__( + stage_name + ) + if stage_values is None: + stage_values = StagePositionValues() + + # Check if the step has already been registered in Murfey + milling_step_entry = murfey_db.exec( + select(MurfeyDB.MillingStep) + .where(MurfeyDB.MillingStep.grid_square_id == grid_square.id) + .where(MurfeyDB.MillingStep.recipe_name == step_info.recipe_name) + .where(MurfeyDB.MillingStep.activity_name == step_info.step_name) + ).one_or_none() + + if milling_step_entry is None: + # Create a new ISPyB entry if no Murfey one is found + result = transport_object.do_insert_milling_step( + # IDs + recipe_name=step_info.recipe_name, + activity_name=step_info.step_name, + grid_square_id=grid_square.id, + # Values + is_enabled=step_info.is_enabled, + status=step_info.status, + execution_time=step_info.execution_time, + stage_x=stage_values.x, + stage_y=stage_values.y, + stage_z=stage_values.z, + rotation=stage_values.rotation, + tilt_alpha=stage_values.tilt_alpha, + beam_type=step_info.beam_type, + beam_voltage=step_info.voltage, + beam_current=step_info.current, + milling_angle=step_info.milling_angle, + depth_correction=step_info.depth_correction, + lamella_offset=step_info.lamella_offset, + trench_height_front=step_info.trench_height_front, + trench_height_rear=step_info.trench_height_rear, + width_overlap_front_left=step_info.width_overlap_front_left, + width_overlap_front_right=step_info.width_overlap_front_right, + width_overlap_rear_left=step_info.width_overlap_rear_left, + width_overlap_rear_right=step_info.width_overlap_rear_right, + ) + if result.get("return_value") is None: + logger.error( + f"No MillingStep entry created for {step_info.step_name}" + ) + continue + + # Create a corresponding record in Murfey + milling_step_entry = MurfeyDB.MillingStep( + id=int(result["return_value"]), + grid_square_id=grid_square.id, + recipe_name=step_info.recipe_name, + activity_name=step_info.step_name, + is_enabled=step_info.is_enabled, + status=step_info.status, + execution_time=step_info.execution_time, + stage_x=stage_values.x, + stage_y=stage_values.y, + stage_z=stage_values.z, + rotation=stage_values.rotation, + tilt_alpha=stage_values.tilt_alpha, + beam_type=step_info.beam_type, + beam_voltage=step_info.voltage, + beam_current=step_info.current, + milling_angle=step_info.milling_angle, + depth_correction=step_info.depth_correction, + lamella_offset=step_info.lamella_offset, + trench_height_front=step_info.trench_height_front, + trench_height_rear=step_info.trench_height_rear, + width_overlap_front_left=step_info.width_overlap_front_left, + width_overlap_front_right=step_info.width_overlap_front_right, + width_overlap_rear_left=step_info.width_overlap_rear_left, + width_overlap_rear_right=step_info.width_overlap_rear_right, + ) + else: + # Update the existing ISPyB one if it already exists + result = transport_object.do_update_milling_step( + milling_step_id=milling_step_entry.id, + is_enabled=step_info.is_enabled, + status=step_info.status, + execution_time=step_info.execution_time, + stage_x=stage_values.x, + stage_y=stage_values.y, + stage_z=stage_values.z, + rotation=stage_values.rotation, + tilt_alpha=stage_values.tilt_alpha, + beam_type=step_info.beam_type, + beam_voltage=step_info.voltage, + beam_current=step_info.current, + milling_angle=step_info.milling_angle, + depth_correction=step_info.depth_correction, + lamella_offset=step_info.lamella_offset, + trench_height_front=step_info.trench_height_front, + trench_height_rear=step_info.trench_height_rear, + width_overlap_front_left=step_info.width_overlap_front_left, + width_overlap_front_right=step_info.width_overlap_front_right, + width_overlap_rear_left=step_info.width_overlap_rear_left, + width_overlap_rear_right=step_info.width_overlap_rear_right, + ) + if result.get("return_value", None) is None: + logger.error( + f"Could not update MillingStep entry for {step_info.step_name}" + ) + continue + + # Update the existing Murfey one + milling_step_entry.is_enabled = step_info.is_enabled + milling_step_entry.status = step_info.status + milling_step_entry.execution_time = step_info.execution_time + milling_step_entry.stage_x = stage_values.x + milling_step_entry.stage_y = stage_values.y + milling_step_entry.stage_z = stage_values.z + milling_step_entry.rotation = stage_values.rotation + milling_step_entry.tilt_alpha = stage_values.tilt_alpha + milling_step_entry.beam_type = step_info.beam_type + milling_step_entry.beam_voltage = step_info.voltage + milling_step_entry.beam_current = step_info.current + milling_step_entry.milling_angle = step_info.milling_angle + milling_step_entry.depth_correction = step_info.depth_correction + milling_step_entry.lamella_offset = step_info.lamella_offset + milling_step_entry.trench_height_front = step_info.trench_height_front + milling_step_entry.trench_height_rear = step_info.trench_height_rear + milling_step_entry.width_overlap_front_left = ( + step_info.width_overlap_front_left + ) + milling_step_entry.width_overlap_front_right = ( + step_info.width_overlap_front_right + ) + milling_step_entry.width_overlap_rear_left = ( + step_info.width_overlap_rear_left + ) + milling_step_entry.width_overlap_rear_right = ( + step_info.width_overlap_rear_right + ) + # Mark entry for committing + murfey_db.add(milling_step_entry) + + # Commit all changes at once + murfey_db.commit() + return None + + +def run(message: dict[str, Any], murfey_db: SQLModelSession): + # Early exit if no TransportManager was set up + if _transport_object is None: + logger.error("No TransportManager object was configured") + return {"success": False, "requeue": False} + + try: + # Parse and unpack incoming message + session_id = int(message["session_id"]) + site_info = LamellaSiteInfo(**message["site_info"]) + logger.debug( + "Received the following FIB metadata for registration:\n" + f"{json.dumps(site_info.model_dump(exclude_none=True), indent=2, default=str)}" + ) + except Exception: + logger.error("Error parsing contents of message", exc_info=True) + return {"success": False, "requeue": False} + + # Early exits if information needed to construct lookup tags are missing + # Project and site values + if site_info.project_name is None: + logger.error("Could not construct lookup tags; 'project_name' is missing") + return {"success": False, "requeue": False} + project_name = site_info.project_name + if site_info.site_number is None: + logger.error("Could not construct lookup tags; 'site_number' is missing") + return {"success": False, "requeue": False} + site_number = site_info.site_number + if site_info.site_name is None: + logger.error("Could not construct lookup tags; 'site_name' is missing") + return {"success": False, "requeue": False} + site_name = site_info.site_name + + # Stage information + if site_info.stage_info is None: + logger.error("Could not construct lookup tags; 'stage_info' is missing") + return {"success": False, "requeue": False} + stage_info = site_info.stage_info + if stage_info.preparation_site is None: + logger.error("Could not construct lookup tags; 'preparation_site' is missing") + return {"success": False, "requeue": False} + preparation_site = stage_info.preparation_site + if preparation_site.slot_number is None: + logger.error("Could not construct lookup tags; 'slot_number' is missing") + return {"success": False, "requeue": False} + slot_number = preparation_site.slot_number + + # Milling step information + if site_info.steps is None: + logger.error("No milling step info found in current message") + return {"success": False, "requeue": False} + milling_steps = site_info.steps + + # Outer try-finally block to handle database cleanup + try: + try: + # Load instrument name and visit ID + murfey_session = murfey_db.exec( + select(MurfeyDB.Session).where(MurfeyDB.Session.id == 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: + # Register the prerequisite information for this site + grid_square_entry = _ensure_prerequisites( + session_id=session_id, + instrument_name=instrument_name, + visit_name=visit_name, + project_name=project_name, + slot_number=slot_number, + site_number=site_number, + transport_object=_transport_object, + murfey_db=murfey_db, + ) + except Exception: + logger.error( + "Exception encountered while registering preqrequisite database entries", + exc_info=True, + ) + return {"success": False, "requeue": False} + if grid_square_entry is None: + logger.error( + f"Could not create GridSquare database entry for site {site_name}" + ) + return {"success": False, "requeue": False} + + try: + # Insert or update MillingStep entries + _register_milling_step( + milling_steps=milling_steps, + stage_info=stage_info, + grid_square=grid_square_entry, + transport_object=_transport_object, + murfey_db=murfey_db, + ) + except Exception: + logger.error( + "Exception encountered while registering milling progress", + exc_info=True, + ) + return {"success": False, "requeue": False} + logger.info(f"Successfully registered milling progress of site {site_name}") + return {"success": True} + finally: + murfey_db.close() diff --git a/tests/server/api/test_workflow_fib.py b/tests/server/api/test_workflow_fib.py index 28b37ae29..1fd400c4d 100644 --- a/tests/server/api/test_workflow_fib.py +++ b/tests/server/api/test_workflow_fib.py @@ -64,7 +64,7 @@ def test_register_fib_atlas( }, ) else: - mock_logger.error.assert_called_with("No Transport Manager object was set up") + mock_logger.error.assert_called_with("No TransportManager object was set up") @pytest.mark.asyncio diff --git a/tests/workflows/fib/test_register_milling_progress.py b/tests/workflows/fib/test_register_milling_progress.py new file mode 100644 index 000000000..606f5ba79 --- /dev/null +++ b/tests/workflows/fib/test_register_milling_progress.py @@ -0,0 +1,539 @@ +from pathlib import Path +from unittest.mock import MagicMock + +import ispyb.sqlalchemy._auto_db_schema as ISPyBDB +import pytest +from pytest_mock import MockerFixture +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_milling_progress import run +from tests.conftest import ExampleVisit + +# Module-wide variables +session_id = 10 +visit_name = f"{ExampleVisit.proposal_code}{ExampleVisit.proposal_number}-{ExampleVisit.visit_number}" +instrument_name = ExampleVisit.instrument_name + + +# Construct test FIB AutoTEM site info for reading (copied and adapted from actual session) +milling_steps = { + "eucentric_tilt": { + "step_name": "Eucentric Tilt", + "recipe_name": "Preparation", + "is_enabled": False, + "status": "Finished", + "execution_time": 328.6657458, + }, + "artificial_features": { + "step_name": "Artificial Features", + "recipe_name": "Preparation", + "is_enabled": False, + "status": "None", + "execution_time": 0.0, + "beam_type": "Ion", + "voltage": 30000.0, + "current": 5e-10, + }, + "milling_angle": { + "step_name": "Milling Angle", + "recipe_name": "Preparation", + "is_enabled": False, + "status": "Finished", + "execution_time": 208.4942922, + "milling_angle": 12.0, + }, + "image_acquisition": { + "step_name": "Image Acquisition", + "recipe_name": "Preparation", + "is_enabled": False, + "status": "Finished", + "execution_time": 19.7872887, + "site_location_type": "Chunk", + }, + "lamella_placement": { + "step_name": "Lamella Placement", + "recipe_name": "Preparation", + "is_enabled": False, + "status": "None", + "execution_time": 0.0, + }, + "delay_1": { + "step_name": "Delay", + "recipe_name": "Milling", + "is_enabled": False, + "status": "None", + "execution_time": 0.0, + }, + "reference_definition": { + "step_name": "Reference Definition", + "recipe_name": "Milling", + "is_enabled": False, + "status": "Finished", + "execution_time": 75.6320442, + "site_location_type": "Chunk", + }, + "reference_definition_electron": { + "step_name": "Electron Reference Definition", + "recipe_name": "Milling", + "is_enabled": False, + "status": "Finished", + "execution_time": 131.7462301, + }, + "stress_relief_cuts": { + "step_name": "Stress Relief Cuts", + "recipe_name": "Milling", + "is_enabled": False, + "status": "None", + "execution_time": 0.0, + "beam_type": "Ion", + "voltage": 30000.0, + "current": 1e-09, + "depth_correction": 3.0, + }, + "reference_redefinition_1": { + "step_name": "Reference Redefinition 1", + "recipe_name": "Milling", + "is_enabled": False, + "status": "None", + "execution_time": 0.0, + "site_location_type": "Chunk", + }, + "rough_milling": { + "step_name": "Rough Milling", + "recipe_name": "Milling", + "is_enabled": False, + "status": "Finished", + "execution_time": 1.5929074719, + "beam_type": "Ion", + "voltage": 30000.0, + "current": 1e-09, + "depth_correction": 3.0, + "lamella_offset": 2e-06, + "trench_height_front": 5.6338138462765e-06, + "trench_height_rear": 8.60473362403628e-06, + "width_overlap_front_left": 2e-06, + "width_overlap_front_right": 2e-06, + "width_overlap_rear_left": 2e-06, + "width_overlap_rear_right": 2e-06, + }, + "rough_milling_electron": { + "step_name": "Rough Milling - Electron Image", + "recipe_name": "Milling", + "is_enabled": False, + "status": "None", + "execution_time": 0.0, + "site_location_type": "Chunk", + "beam_type": "Electron", + "voltage": 2000.0, + "current": 1.25e-11, + }, + "reference_redefinition_2": { + "step_name": "Reference Redefinition 2", + "recipe_name": "Milling", + "is_enabled": False, + "status": "Finished", + "execution_time": 105.1459769, + "site_location_type": "Chunk", + }, + "medium_milling": { + "step_name": "Medium Milling", + "recipe_name": "Milling", + "is_enabled": False, + "status": "Finished", + "execution_time": 298.1471377, + "beam_type": "Ion", + "voltage": 30000.0, + "current": 5e-10, + "depth_correction": 2.0, + "lamella_offset": 1.5e-06, + "width_overlap_front_left": 1.5e-06, + "width_overlap_front_right": 1.5e-06, + "width_overlap_rear_left": 1.5e-06, + "width_overlap_rear_right": 1.5e-06, + }, + "medium_milling_electron": { + "step_name": "Medium Milling - Electron Image", + "recipe_name": "Milling", + "is_enabled": False, + "status": "None", + "execution_time": 0.0, + "site_location_type": "Chunk", + "beam_type": "Electron", + "voltage": 2000.0, + "current": 1.25e-11, + }, + "fine_milling": { + "step_name": "Fine Milling", + "recipe_name": "Milling", + "is_enabled": False, + "status": "Finished", + "execution_time": 397.414388, + "beam_type": "Ion", + "voltage": 30000.0, + "current": 3e-10, + "depth_correction": 2.0, + "lamella_offset": 1e-06, + "width_overlap_front_left": 1e-06, + "width_overlap_front_right": 1e-06, + "width_overlap_rear_left": 1e-06, + "width_overlap_rear_right": 1e-06, + }, + "fine_milling_electron": { + "step_name": "Fine Milling - Electron Image", + "recipe_name": "Milling", + "is_enabled": False, + "status": "None", + "execution_time": 0.0, + "site_location_type": "Chunk", + "beam_type": "Electron", + "voltage": 2000.0, + "current": 1.25e-11, + }, + "finer_milling": { + "step_name": "Finer Milling", + "recipe_name": "Milling", + "is_enabled": False, + "status": "Finished", + "execution_time": 887.9686893, + "beam_type": "Ion", + "voltage": 30000.0, + "current": 1e-10, + "depth_correction": 2.0, + "lamella_offset": 4.0000000000000003e-07, + "width_overlap_front_left": 5.000000000000001e-07, + "width_overlap_front_right": 5.000000000000001e-07, + "width_overlap_rear_left": 5.000000000000001e-07, + "width_overlap_rear_right": 5.000000000000001e-07, + }, + "finer_milling_electron": { + "step_name": "Finer Milling - Electron Image", + "recipe_name": "Milling", + "is_enabled": False, + "status": "Finished", + "execution_time": 59.2296577, + "site_location_type": "Chunk", + "beam_type": "Electron", + "voltage": 2000.0, + "current": 2.5e-11, + }, + "delay_2": { + "step_name": "Delay", + "recipe_name": "Thinning", + "is_enabled": False, + "status": "None", + "execution_time": 0.0, + }, + "polishing_1": { + "step_name": "Polishing 1", + "recipe_name": "Thinning", + "is_enabled": False, + "status": "Finished", + "execution_time": 680.5567315, + "beam_type": "Ion", + "voltage": 30000.0, + "current": 5e-11, + "depth_correction": 2.0, + "lamella_offset": 2.5000000000000004e-07, + "width_overlap_front_left": 0.0, + "width_overlap_front_right": 0.0, + "width_overlap_rear_left": 0.0, + "width_overlap_rear_right": 0.0, + }, + "polishing_1_electron": { + "step_name": "Polishing 1 - Electron Image", + "recipe_name": "Thinning", + "is_enabled": False, + "status": "None", + "execution_time": 0.0, + "site_location_type": "Thinning", + "beam_type": "Electron", + "voltage": 2000.0, + "current": 1.25e-11, + }, + "polishing_2": { + "step_name": "Polishing 2", + "recipe_name": "Thinning", + "is_enabled": False, + "status": "Finished", + "execution_time": 1.170488927, + "beam_type": "Ion", + "voltage": 30000.0, + "current": 3e-11, + "depth_correction": 1.5, + "lamella_offset": 0.0, + "width_overlap_front_left": 0.0, + "width_overlap_front_right": 0.0, + "width_overlap_rear_left": 0.0, + "width_overlap_rear_right": 0.0, + }, + "polishing_2_ion": { + "step_name": "Polishing 2 - Ion Image", + "recipe_name": "Thinning", + "is_enabled": False, + "status": "None", + "execution_time": 0.0, + "site_location_type": "Thinning", + "beam_type": "Ion", + "voltage": 30000.0, + "current": 3e-11, + }, + "polishing_2_electron": { + "step_name": "Polishing 2 - Electron Image", + "recipe_name": "Thinning", + "is_enabled": False, + "status": "Finished", + "execution_time": 56.9180832, + "site_location_type": "Thinning", + "beam_type": "Electron", + "voltage": 2000.0, + "current": 2.5e-11, + }, +} +stage_info = { + "preparation_site": { + "x": -0.0031091224743875403, + "y": 0.00420867925495798, + "z": 0.0323644854106331, + "rotation": 285.003247202109, + "tilt_alpha": 25.9996646026832, + "slot_number": 1, + }, + "chunk_site": { + "x": -0.0030037500000000003, + "y": 0.004293, + "z": 0.032350405092592606, + "rotation": 285.003247202109, + "tilt_alpha": -0.000134158926728586, + "slot_number": 1, + }, + "thinning_site": { + "x": -0.0030037500000000003, + "y": 0.004293, + "z": 0.032350405092592606, + "rotation": 285.003247202109, + "tilt_alpha": -0.000134158926728586, + "slot_number": 1, + }, + "chunk_coincidence_params": { + "x": -0.0030048260286678298, + "y": 0.004308828126160981, + "z": 0.0323400707790533, + "rotation": 285.003247202109, + "tilt_alpha": -0.000134158926728586, + "slot_number": 1, + }, + "thinning_params": { + "x": -0.0030037500000000003, + "y": 0.004293, + "z": 0.032350405092592606, + "rotation": 285.003247202109, + "tilt_alpha": -0.000134158926728586, + "slot_number": 1, + }, +} +site_info = { + "project_name": visit_name, + "site_name": "Lamella", + "site_number": 1, + "stage_info": stage_info, + "steps": milling_steps, +} +# Construct the RabbitMQ message received +message = { + "register": "fib.register_milling_progress", + "session_id": session_id, + "site_info": site_info, +} + + +def test_run_with_db( + mocker: MockerFixture, + murfey_db_session: SQLModelSession, + ispyb_db_session: SQLAlchemySession, + mock_ispyb_credentials: Path, +): + # Add a test visit to the database + if not ( + session_entry := 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_db_session.add(session_entry) + murfey_db_session.commit() + + # 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 DataCollectionGroup + 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.fib.register_milling_progress._transport_object", + new=TransportManager("PikaTransport"), + ) + + # Run the workflow twice (fresh insert and update existing) + for i in range(2): + result = run( + message=message, + murfey_db=murfey_db_session, + ) + assert result.get("success", False) + + # There should be a DataCollectionGroup entry in Murfey + dcg_murfey = murfey_db_session.exec( + sm_select(MurfeyDB.DataCollectionGroup) + .where(MurfeyDB.DataCollectionGroup.session_id == session_id) + .where( + MurfeyDB.DataCollectionGroup.tag == f"{site_info['project_name']}--slot_1" + ) + ).one_or_none() + assert dcg_murfey is not None + + # There should be a DataCollectionGroup entry in ISPyB + dcg_ispyb = ispyb_db_session.execute( + sa_select(ISPyBDB.DataCollectionGroup).where( + ISPyBDB.DataCollectionGroup.dataCollectionGroupId == dcg_murfey.id + ) + ).scalar_one_or_none() + assert dcg_ispyb is not None + + # There should be an Atlas in ISPyB + atlas_ispyb = ispyb_db_session.execute( + sa_select(ISPyBDB.Atlas).where( + ISPyBDB.Atlas.dataCollectionGroupId == dcg_ispyb.dataCollectionGroupId + ) + ).scalar_one_or_none() + assert atlas_ispyb is not None + + # There should be one GridSquare entry in ISPyB per lamella site tested + gs_ispyb_rows = ( + ispyb_db_session.execute( + sa_select(ISPyBDB.GridSquare).where( + ISPyBDB.GridSquare.atlasId == atlas_ispyb.atlasId + ) + ) + .scalars() + .all() + ) + assert len(gs_ispyb_rows) >= 1 + gs_ispyb = gs_ispyb_rows[0] + + # There should be one MillingStep entry in ISPyB for each step in the message + milling_step_ispyb_rows = ( + ispyb_db_session.execute( + sa_select(ISPyBDB.MillingStep).where( + ISPyBDB.MillingStep.gridSquareId == gs_ispyb.gridSquareId + ) + ) + .scalars() + .all() + ) + assert len(milling_step_ispyb_rows) == len(milling_steps) + + # There should be the same thing in Murfey + milling_step_murfey_rows = murfey_db_session.exec( + sm_select(MurfeyDB.MillingStep).where( + MurfeyDB.MillingStep.grid_square_id == gs_ispyb.gridSquareId + ) + ).all() + assert len(milling_step_murfey_rows) == len(milling_steps) + + +@pytest.mark.parametrize( + "test_params", + ( # TransportManager | Project name | Site number | Site name | Stage info | Preparation site | Slot number + (False, False, False, False, False, False, False), + (True, False, False, False, False, False, False), + (True, True, False, False, False, False, False), + (True, True, True, False, False, False, False), + (True, True, True, True, False, False, False), + (True, True, True, True, True, False, False), + (True, True, True, True, True, True, False), + (True, True, True, True, True, True, True), + ), +) +def test_run_fails( + mocker: MockerFixture, + test_params: tuple[bool, bool, bool, bool, bool, bool, bool], +): + # Unpack test params + ( + has_transport_object, + has_project_name, + has_site_number, + has_site_name, + has_stage_info, + has_preparation_site, + has_slot_number, + ) = test_params + + # Mock TransportManager accordingly + mock_transport_object = MagicMock() if has_transport_object else None + mocker.patch( + "murfey.workflows.fib.register_milling_progress._transport_object", + mock_transport_object, + ) + + # Mock the Murfey database + mock_murfey_db = MagicMock() + + # Mock the site info dictionary accordingly + project_name = visit_name if has_project_name else None + site_name = "Lamella" if has_site_name else None + site_num = 1 if has_site_number else None + message = { + "register": "fib.register_milling_progress", + "session_id": session_id, + "site_info": { + "project_name": project_name, + "site_name": site_name, + "site_number": site_num, + "stage_info": { + "preparation_site": { + "x": -0.0031091224743875403 if has_slot_number else None, + "y": 0.00420867925495798, + "z": 0.0323644854106331, + "rotation": 285.003247202109, + "tilt_alpha": 25.9996646026832, + "slot_number": 1, + } + if has_preparation_site + else None, + } + if has_stage_info + else None, + "steps": None, + }, + } + + # Run the function and check that the correct message is returned + result = run(message, mock_murfey_db) + assert result == {"success": False, "requeue": False}