Move the spa flush to a workflow

Author: stephen-riggs

pyproject.toml

@@ -107,6 +107,7 @@ murfey = "murfey.client:run"
 "clem.register_align_and_merge_result" = "murfey.workflows.clem.register_align_and_merge_results:register_align_and_merge_result"
 "clem.register_lif_preprocessing_result" = "murfey.workflows.clem.register_preprocessing_results:register_lif_preprocessing_result"
 "clem.register_tiff_preprocessing_result" = "murfey.workflows.clem.register_preprocessing_results:register_tiff_preprocessing_result"
+"spa.flush_spa_preprocess" = "murfey.workflows.spa.flush_spa_preprocess:flush_spa_preprocess"
 
 [tool.setuptools]
 package-dir = {"" = "src"}

src/murfey/server/__init__.py

@@ -11,7 +11,7 @@
 from importlib.resources import files
 from pathlib import Path
 from threading import Thread
-from typing import Any, Dict, List, Literal, NamedTuple, Optional, Tuple
+from typing import Any, Dict, List, Literal, NamedTuple, Tuple
 
 import graypy
 import mrcfile
@@ -49,7 +49,6 @@
 import murfey.server.websocket
 import murfey.util.db as db
 from murfey.server.murfey_db import url  # murfey_db
-from murfey.server.spa.spa_metadata import register_foil_hole, register_grid_square
 from murfey.util import LogFilter
 from murfey.util.config import (
     MachineConfig,
@@ -58,16 +57,7 @@
     get_microscope,
     get_security_config,
 )
-from murfey.util.models import FoilHoleParameters, GridSquareParameters
 from murfey.util.processing_params import default_spa_parameters
-from murfey.util.spa_metadata import (
-    GridSquare,
-    _foil_hole_data,
-    _foil_hole_from_file,
-    _get_grid_square_atlas_positions,
-    _grid_square_data,
-    _grid_square_from_file,
-)
 from murfey.util.state import global_state
 from murfey.util.tomo import midpoint
 
@@ -1940,223 +1930,6 @@ def _register_initial_model(message: dict, _db=murfey_db, demo: bool = False):
     _db.close()
 
 
-def _grid_square_metadata_file(f: Path, grid_square: int) -> Optional[Path]:
-    """Search through metadata directories to find the required grid square dm"""
-    raw_dir = f.parent.parent.parent
-    metadata_dirs = raw_dir.glob("metadata*")
-    for md_dir in metadata_dirs:
-        gs_path = md_dir / f"Metadata/GridSquare_{grid_square}.dm"
-        if gs_path.is_file():
-            return gs_path
-    logger.error(f"Could not determine grid square metadata path for {f}")
-    return None
-
-
-def _flush_position_analysis(
-    movie_path: Path, dcg_id: int, session_id: int
-) -> Optional[int]:
-    """Register a grid square and foil hole in the database"""
-    data_collection_group = murfey_db.exec(
-        select(db.DataCollectionGroup).where(db.DataCollectionGroup.id == dcg_id)
-    ).one()
-
-    # Work out the grid square and associated metadata file
-    grid_square = _grid_square_from_file(movie_path)
-    grid_square_metadata_file = _grid_square_metadata_file(movie_path, grid_square)
-    if grid_square_metadata_file:
-        gs = _grid_square_data(grid_square_metadata_file, grid_square)
-    else:
-        gs = GridSquare(id=grid_square)
-    if data_collection_group.atlas:
-        # If an atlas if present, work out where this grid square is on it
-        gs_pix_position = _get_grid_square_atlas_positions(
-            data_collection_group.atlas,
-            grid_square=str(grid_square),
-        )[str(grid_square)]
-        grid_square_parameters = GridSquareParameters(
-            tag=data_collection_group.tag,
-            x_location=gs_pix_position[0],
-            y_location=gs_pix_position[1],
-            x_stage_position=gs_pix_position[2],
-            y_stage_position=gs_pix_position[3],
-            readout_area_x=gs.readout_area_x,
-            readout_area_y=gs.readout_area_y,
-            thumbnail_size_x=gs.thumbnail_size_x,
-            thumbnail_size_y=gs.thumbnail_size_y,
-            height=gs_pix_position[5],
-            width=gs_pix_position[4],
-            pixel_size=gs.pixel_size,
-            image=gs.image,
-            angle=gs_pix_position[6],
-        )
-    else:
-        # Skip location analysis if no atlas
-        grid_square_parameters = GridSquareParameters(
-            tag=data_collection_group.tag,
-            readout_area_x=gs.readout_area_x,
-            readout_area_y=gs.readout_area_y,
-            thumbnail_size_x=gs.thumbnail_size_x,
-            thumbnail_size_y=gs.thumbnail_size_y,
-            pixel_size=gs.pixel_size,
-            image=gs.image,
-        )
-    # Insert or update this grid square in the database
-    register_grid_square(session_id, gs.id, grid_square_parameters, murfey_db)
-
-    # Find the foil hole info and register it
-    foil_hole = _foil_hole_from_file(movie_path)
-    if grid_square_metadata_file:
-        fh = _foil_hole_data(
-            grid_square_metadata_file,
-            foil_hole,
-            grid_square,
-        )
-        foil_hole_parameters = FoilHoleParameters(
-            tag=data_collection_group.tag,
-            name=foil_hole,
-            x_location=fh.x_location,
-            y_location=fh.y_location,
-            x_stage_position=fh.x_stage_position,
-            y_stage_position=fh.y_stage_position,
-            readout_area_x=fh.readout_area_x,
-            readout_area_y=fh.readout_area_y,
-            thumbnail_size_x=fh.thumbnail_size_x,
-            thumbnail_size_y=fh.thumbnail_size_y,
-            pixel_size=fh.pixel_size,
-            image=fh.image,
-            diameter=fh.diameter,
-        )
-    else:
-        foil_hole_parameters = FoilHoleParameters(
-            tag=data_collection_group.tag,
-            name=foil_hole,
-        )
-    # Insert or update this foil hole in the database
-    register_foil_hole(session_id, gs.id, foil_hole_parameters, murfey_db)
-    return foil_hole
-
-
-def _flush_spa_preprocessing(message: dict):
-    session_id = message["session_id"]
-    stashed_files = murfey_db.exec(
-        select(db.PreprocessStash)
-        .where(db.PreprocessStash.session_id == session_id)
-        .where(db.PreprocessStash.tag == message["tag"])
-    ).all()
-    if not stashed_files:
-        return None
-    instrument_name = (
-        murfey_db.exec(select(db.Session).where(db.Session.id == message["session_id"]))
-        .one()
-        .instrument_name
-    )
-    machine_config = get_machine_config(instrument_name=instrument_name)[
-        instrument_name
-    ]
-    recipe_name = machine_config.recipes.get("em-spa-preprocess", "em-spa-preprocess")
-    collected_ids = murfey_db.exec(
-        select(
-            db.DataCollectionGroup,
-            db.DataCollection,
-            db.ProcessingJob,
-            db.AutoProcProgram,
-        )
-        .where(db.DataCollectionGroup.session_id == session_id)
-        .where(db.DataCollectionGroup.tag == message["tag"])
-        .where(db.DataCollection.dcg_id == db.DataCollectionGroup.id)
-        .where(db.ProcessingJob.dc_id == db.DataCollection.id)
-        .where(db.AutoProcProgram.pj_id == db.ProcessingJob.id)
-        .where(db.ProcessingJob.recipe == recipe_name)
-    ).one()
-    params = murfey_db.exec(
-        select(db.SPARelionParameters, db.SPAFeedbackParameters)
-        .where(db.SPARelionParameters.pj_id == collected_ids[2].id)
-        .where(db.SPAFeedbackParameters.pj_id == db.SPARelionParameters.pj_id)
-    ).one()
-    proc_params = params[0]
-    feedback_params = params[1]
-    if not proc_params:
-        logger.warning(
-            f"No SPA processing parameters found for client processing job ID {collected_ids[2].id}"
-        )
-        raise ValueError(
-            "No processing parameters were found in the database when flushing SPA preprocessing"
-        )
-
-    murfey_ids = _murfey_id(
-        collected_ids[3].id,
-        murfey_db,
-        number=2 * len(stashed_files),
-        close=False,
-    )
-    if feedback_params.picker_murfey_id is None:
-        feedback_params.picker_murfey_id = murfey_ids[1]
-        murfey_db.add(feedback_params)
-
-    for i, f in enumerate(stashed_files):
-        if f.foil_hole_id:
-            foil_hole_id = f.foil_hole_id
-        else:
-            # Register grid square and foil hole if not present
-            foil_hole_id = _flush_position_analysis(
-                movie_path=f.file_path,
-                dcg_id=collected_ids[0].id,
-                session_id=session_id,
-            )
-
-        mrcp = Path(f.mrc_out)
-        ppath = Path(f.file_path)
-        if not mrcp.parent.exists():
-            mrcp.parent.mkdir(parents=True)
-        movie = db.Movie(
-            murfey_id=murfey_ids[2 * i],
-            path=f.file_path,
-            image_number=f.image_number,
-            tag=f.tag,
-            foil_hole_id=foil_hole_id,
-        )
-        murfey_db.add(movie)
-        zocalo_message: dict = {
-            "recipes": [recipe_name],
-            "parameters": {
-                "node_creator_queue": machine_config.node_creator_queue,
-                "dcid": collected_ids[1].id,
-                "kv": proc_params.voltage,
-                "autoproc_program_id": collected_ids[3].id,
-                "movie": f.file_path,
-                "mrc_out": f.mrc_out,
-                "pixel_size": proc_params.angpix,
-                "image_number": f.image_number,
-                "microscope": get_microscope(),
-                "mc_uuid": murfey_ids[2 * i],
-                "ft_bin": proc_params.motion_corr_binning,
-                "fm_dose": proc_params.dose_per_frame,
-                "gain_ref": proc_params.gain_ref,
-                "picker_uuid": murfey_ids[2 * i + 1],
-                "session_id": session_id,
-                "particle_diameter": proc_params.particle_diameter or 0,
-                "fm_int_file": f.eer_fractionation_file,
-                "do_icebreaker_jobs": default_spa_parameters.do_icebreaker_jobs,
-                "foil_hole_id": foil_hole_id,
-            },
-        }
-        if _transport_object:
-            zocalo_message["parameters"][
-                "feedback_queue"
-            ] = _transport_object.feedback_queue
-            _transport_object.send(
-                "processing_recipe", zocalo_message, new_connection=True
-            )
-            murfey_db.delete(f)
-        else:
-            logger.error(
-                f"Pre-processing was requested for {ppath.name} but no Zocalo transport object was found"
-            )
-    murfey_db.commit()
-    murfey_db.close()
-    return None
-
-
 def _flush_tomography_preprocessing(message: dict):
     session_id = message["session_id"]
     instrument_name = (
@@ -2956,11 +2729,6 @@ def feedback_callback(header: dict, message: dict) -> None:
             if _transport_object:
                 _transport_object.transport.ack(header)
             return None
-        elif message["register"] == "flush_spa_preprocess":
-            _flush_spa_preprocessing(message)
-            if _transport_object:
-                _transport_object.transport.ack(header)
-            return None
         elif message["register"] == "spa_processing_parameters":
             session_id = message["session_id"]
             collected_ids = murfey_db.exec(