SmartEM integration for the EPU context (#772)

Add a smartem grid when a data collection group is made

Add smartem gridsquares alongside Murfey squares, and the same for foilholes

Once all squares have been added register the grid in smartem to trigger quality predictions

Similar for foilholes on each square

Add some post processing stage hooks to pass the correct measurement information to smartem

Not fully tested for robustness against different metadata file sequences but designed to not interfere with current capabilities

Author: d-j-hatton

pyproject.toml

@@ -119,7 +119,9 @@ TomographyMetadataContext = "murfey.client.contexts.tomo_metadata:TomographyMeta
 "picked_particles" = "murfey.workflows.spa.picking:particles_picked"
 "picked_tomogram" = "murfey.workflows.tomo.picking:picked_tomogram"
 "processing_job" = "murfey.workflows.register_processing_job:run"
+"spa.ctf_estimated" = "murfey.workflows.spa.ctf_estimation:ctf_estimated"
 "spa.flush_spa_preprocess" = "murfey.workflows.spa.flush_spa_preprocess:flush_spa_preprocess"
+"spa.motion_corrected" = "murfey.workflows.spa.motion_correction:motion_corrected"
 
 [tool.setuptools]
 package-dir = {"" = "src"}

src/murfey/client/analyser.py

@@ -303,7 +303,7 @@ def _find_context(self, file_path: Path) -> bool:
 
             # Files starting with "Position" belong to the standard tomography workflow
             # NOTE: not completely reliable, mdocs can be in tomography metadata as well
-            if (
+            if not self._serialem and (
                 split_file_stem[0] == "Position"
                 or "[" in file_path.name
                 or split_file_stem[-1] in ["Fractions", "fractions", "EER"]
@@ -378,7 +378,7 @@ def _analyse(self):
                 self.post_transfer(transferred_file)
             else:
                 dc_metadata = {}
-                if (
+                if not self._serialem and (
                     self._force_mdoc_metadata
                     and transferred_file.suffix == ".mdoc"
                     or mdoc_for_reading

src/murfey/client/contexts/atlas.py

@@ -60,6 +60,7 @@ def post_transfer_serialem(
                     data={
                         "name": transferred_file.stem,
                         "acquisition_uuid": environment.acquisition_uuid,
+                        "storage_folder": str(source),
                     },
                 )
 
@@ -134,6 +135,8 @@ def post_transfer_epu(
                     "atlas": str(transferred_atlas_jpg).replace("//", "/"),
                     "sample": sample,
                     "atlas_pixel_size": atlas_pixel_size,
+                    "create_smartem_grid": bool(environment.acquisition_uuid),
+                    "acquisition_uuid": environment.acquisition_uuid,
                 }
                 capture_post(
                     base_url=str(environment.url.geturl()),

src/murfey/client/contexts/spa_metadata.py

@@ -152,6 +152,19 @@ def post_transfer(
                                     "angle": pos_data[6],
                                 },
                             )
+                    if pos_data:
+                        capture_post(
+                            base_url=str(environment.url.geturl()),
+                            router_name="session_control.spa_router",
+                            function_name="register_atlas",
+                            token=self._token,
+                            session_id=environment.murfey_session,
+                            data={
+                                "name": f"{environment.visit}-sample-{environment.samples[images_disc].sample}",
+                                "acquisition_uuid": environment.acquisition_uuid,
+                                "register_grid": True,
+                            },
+                        )
 
         elif (
             transferred_file.suffix == ".dm"
@@ -253,3 +266,13 @@ def post_transfer(
                         "image": fh_data.image,
                     },
                 )
+            if fh_positions:
+                capture_post(
+                    base_url=str(environment.url.geturl()),
+                    router_name="session_control.spa_router",
+                    function_name="register_square",
+                    token=self._token,
+                    session_id=environment.murfey_session,
+                    gsid=gs_name,
+                    data={"tag": visitless_source},
+                )

src/murfey/server/api/session_control.py

@@ -360,6 +360,8 @@ def get_foil_hole(
 class AtlasRegistration(BaseModel):
     name: str
     acquisition_uuid: str
+    storage_folder: str = ""
+    register_grid: bool = False
 
 
 @spa_router.post("/sessions/{session_id}/register_atlas")
@@ -389,17 +391,51 @@ def register_atlas(
                 atlas_data = AtlasData(
                     id=atlas_registration_data.name,
                     acquisition_date=datetime.now(),
-                    storage_folder="",
+                    storage_folder=atlas_registration_data.storage_folder,
                     name=atlas_registration_data.name,
                     tiles=[],
                     gridsquare_positions=None,
                     grid_uuid=grid_uuid,
                 )
                 smartem_client.create_grid_atlas(atlas_data)
+                if atlas_registration_data.register_grid:
+                    smartem_client.grid_registered(grid_uuid)
     else:
         logger.info("smartem deactivated so did not register atlas")
 
 
+class SquareRegistration(BaseModel):
+    tag: str
+
+
+@spa_router.post("/sessions/{session_id}/register_square/{gsid}")
+def register_square(
+    session_id: MurfeySessionID,
+    gsid: int,
+    square_registration_data: SquareRegistration,
+    db=murfey_db,
+):
+    if SMARTEM_ACTIVE:
+        gs = db.exec(
+            select(GridSquare)
+            .where(GridSquare.name == gsid)
+            .where(GridSquare.tag == square_registration_data.tag)
+            .where(GridSquare.session_id == session_id)
+        ).one_or_none()
+        if gs and gs.smartem_uuid:
+            session = db.exec(select(Session).where(Session.id == session_id)).one()
+            machine_config = get_machine_config(session.instrument_name)[
+                session.instrument_name
+            ]
+            if machine_config.smartem_api_url:
+                smartem_client = SmartEMAPIClient(
+                    base_url=machine_config.smartem_api_url, logger=logger
+                )
+                smartem_client.gridsquare_registered(gs.smartem_uuid)
+    else:
+        logger.info("smartem deactivated so did not register square")
+
+
 @spa_router.post("/sessions/{session_id}/make_atlas_jpg")
 def make_atlas_jpg(
     session_id: MurfeySessionID, atlas_mrc: StringOfPathModel, db=murfey_db

src/murfey/server/api/workflow.py

@@ -24,6 +24,19 @@
 except ImportError:
     Image = None
 
+try:
+    from smartem_backend.api_client import SmartEMAPIClient
+    from smartem_common.schemas import (
+        AcquisitionData as SmartEMAcquisitionData,
+        GridData as SmartEMGridData,
+        MicrographData as SmartEMMicrographData,
+        MicrographManifest as SmartEMMicrographManifest,
+    )
+
+    SMARTEM_ACTIVE = True
+except ImportError:
+    SMARTEM_ACTIVE = False
+
 import murfey.server.prometheus as prom
 from murfey.server import _transport_object
 from murfey.server.api.auth import (
@@ -78,6 +91,7 @@
 
 logger = getLogger("murfey.server.api.workflow")
 
+
 router = APIRouter(
     prefix="/workflow",
     dependencies=[Depends(validate_instrument_token)],
@@ -92,6 +106,8 @@ class DCGroupParameters(BaseModel):
     atlas: str = ""
     sample: Optional[int] = None
     atlas_pixel_size: float = 0
+    create_smartem_grid: bool = False
+    acquisition_uuid: Optional[str] = None
 
 
 @router.post(
@@ -110,6 +126,35 @@ def register_dc_group(
         db.exec(select(Session).where(Session.id == session_id)).one().instrument_name
     )
     logger.info(f"Registering data collection group on microscope {instrument_name}")
+    smartem_grid_uuid = None
+    if (
+        dcg_params.create_smartem_grid
+        and SMARTEM_ACTIVE
+        and dcg_params.acquisition_uuid
+    ):
+        machine_config = get_machine_config(instrument_name=instrument_name)[
+            instrument_name
+        ]
+        if machine_config.smartem_api_url:
+            try:
+                smartem_client = SmartEMAPIClient(
+                    base_url=machine_config.smartem_api_url, logger=logger
+                )
+                grid_data = SmartEMGridData(
+                    data_dir=Path(dcg_params.tag),
+                    atlas_dir=Path(dcg_params.atlas) if dcg_params.atlas else None,
+                    acquisition_data=SmartEMAcquisitionData(
+                        uuid=dcg_params.acquisition_uuid,
+                        name=f"{visit_name}-sample-{dcg_params.sample}"
+                        if dcg_params.sample
+                        else f"{visit_name}-sample-unknown",
+                    ),
+                )
+                smartem_grid_uuid = smartem_client.create_acquisition_grid(
+                    grid_data
+                ).uuid
+            except Exception:
+                logger.warning("Failed to register SmartEM grid", exc_info=True)
     if (
         dcg_murfey := db.exec(
             select(DataCollectionGroup)
@@ -135,6 +180,8 @@ def register_dc_group(
             dcg_instance.atlas_pixel_size = (
                 dcg_params.atlas_pixel_size or dcg_instance.atlas_pixel_size
             )
+            if smartem_grid_uuid:
+                dcg_instance.smartem_grid_uuid = smartem_grid_uuid
 
             if _transport_object:
                 if dcg_instance.atlas_id is not None:
@@ -217,6 +264,11 @@ def register_dc_group(
                     "proposal_code": ispyb_proposal_code,
                     "proposal_number": ispyb_proposal_number,
                     "visit_number": ispyb_visit_number,
+                    **(
+                        {"smartem_grid_uuid": smartem_grid_uuid}
+                        if smartem_grid_uuid
+                        else {}
+                    ),
                 },
             )
     return dcg_params
@@ -491,6 +543,58 @@ async def request_spa_preprocessing(
         )
         db.add(movie)
         db.commit()
+
+        if (
+            SMARTEM_ACTIVE
+            and machine_config.smartem_api_url
+            and foil_hole_id is not None
+        ):
+            try:
+                fh_with_gs = db.exec(
+                    select(FoilHole, GridSquare)
+                    .where(FoilHole.id == foil_hole_id)
+                    .where(GridSquare.id == FoilHole.grid_square_id)
+                ).one_or_none()
+                if fh_with_gs is not None:
+                    fh, gs = fh_with_gs
+                    if fh.smartem_uuid:
+                        smartem_client = SmartEMAPIClient(
+                            base_url=machine_config.smartem_api_url, logger=logger
+                        )
+                        movie_path = Path(proc_file.path)
+                        micrograph_manifest = SmartEMMicrographManifest(
+                            unique_id=movie_path.stem,
+                            acquisition_datetime=datetime.now(),
+                            defocus=None,
+                            detector_name="",
+                            energy_filter=True,
+                            phase_plate=False,
+                            image_size_x=None,
+                            image_size_y=None,
+                            binning_x=1,
+                            binning_y=1,
+                        )
+                        micrograph_data = SmartEMMicrographData(
+                            id=movie_path.stem,
+                            gridsquare_id=str(gs.name),
+                            foilhole_uuid=fh.smartem_uuid,
+                            foilhole_id=str(fh.name),
+                            location_id=str(murfey_ids[0]),
+                            high_res_path=movie_path,
+                            manifest_file=movie_path,
+                            manifest=micrograph_manifest,
+                        )
+                        response = smartem_client.create_foilhole_micrograph(
+                            micrograph_data
+                        )
+                        movie.smartem_uuid = response.uuid
+                        db.add(movie)
+                        db.commit()
+            except Exception:
+                logger.warning(
+                    "Failed to register micrograph with smartem", exc_info=True
+                )
+
         db.close()
 
         if not mrc_out.parent.exists():

src/murfey/util/db.py

@@ -211,6 +211,7 @@ class DataCollectionGroup(SQLModel, table=True):  # type: ignore
     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",
@@ -500,6 +501,7 @@ class GridSquare(SQLModel, table=True):  # type: ignore
     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"
     )
@@ -533,6 +535,7 @@ class FoilHole(SQLModel, table=True):  # type: ignore
     pixel_location_y: Optional[int] = None
     diameter: Optional[int] = None
     quality_indicator: Optional[float] = None
+    smartem_uuid: Optional[str] = None
 
 
 class SearchMap(SQLModel, table=True):  # type: ignore
@@ -612,6 +615,7 @@ class Movie(SQLModel, table=True):  # type: ignore
     fluence: Optional[float] = None
     numberOfFrames: Optional[int] = None
     templateLabel: Optional[str] = None
+    smartem_uuid: Optional[str] = None
     murfey_ledger: Optional[MurfeyLedger] = Relationship(back_populates="movies")
     data_collection: Optional["DataCollection"] = Relationship(back_populates="movies")
     foil_hole: Optional[FoilHole] = Relationship(back_populates="movies")

src/murfey/util/route_manifest.yaml

@@ -947,6 +947,15 @@ murfey.server.api.session_control.spa_router:
         type: int
     methods:
       - POST
+  - path: /session_control/spa/sessions/{session_id}/register_square/{gsid}
+    function: register_square
+    path_params:
+      - name: gsid
+        type: int
+      - name: session_id
+        type: int
+    methods:
+      - POST
   - path: /session_control/spa/sessions/{session_id}/make_atlas_jpg
     function: make_atlas_jpg
     path_params:

src/murfey/workflows/register_data_collection_group.py

@@ -41,6 +41,7 @@ def run(message: dict, murfey_db: SQLModelSession) -> dict[str, bool]:
             murfey_dcg = DataCollectionGroup(
                 session_id=message["session_id"],
                 tag=message.get("tag"),
+                smartem_grid_uuid=message.get("smartem_grid_uuid"),
             )
             dcgid = murfey_dcg.id
         else:
@@ -86,6 +87,7 @@ def run(message: dict, murfey_db: SQLModelSession) -> dict[str, bool]:
                 sample=message.get("sample"),
                 session_id=message["session_id"],
                 tag=message.get("tag"),
+                smartem_grid_uuid=message.get("smartem_grid_uuid"),
             )
         murfey_db.add(murfey_dcg)
         murfey_db.commit()