Merged recent changes from 'main' branch

Author: tieneupin

src/murfey/client/contexts/atlas.py

@@ -7,6 +7,7 @@
 from murfey.client.context import Context, _atlas_destination, _get_source
 from murfey.client.instance_environment import MurfeyInstanceEnvironment
 from murfey.util.client import capture_post
+from murfey.util.spa_metadata import get_grid_square_atlas_positions
 
 logger = logging.getLogger("murfey.client.contexts.atlas")
 
@@ -158,3 +159,52 @@ def post_transfer_epu(
                 logger.info(
                     f"Registered data collection group for atlas {str(transferred_atlas_jpg)!r}"
                 )
+
+        elif environment and transferred_file.name == "Atlas.dm":
+            # Register all grid squares on this atlas
+            gs_pix_positions = get_grid_square_atlas_positions(transferred_file)
+            for gs, pos_data in gs_pix_positions.items():
+                if pos_data:
+                    capture_post(
+                        base_url=str(environment.url.geturl()),
+                        router_name="session_control.spa_router",
+                        function_name="register_grid_square",
+                        token=self._token,
+                        instrument_name=environment.instrument_name,
+                        session_id=environment.murfey_session,
+                        gsid=int(gs),
+                        data={
+                            "tag": str(transferred_file.parent),
+                            "x_location": pos_data[0],
+                            "y_location": pos_data[1],
+                            "x_stage_position": pos_data[2],
+                            "y_stage_position": pos_data[3],
+                            "width": pos_data[4],
+                            "height": pos_data[5],
+                            "angle": pos_data[6],
+                        },
+                    )
+            if gs_pix_positions:
+                for p in transferred_file.parts:
+                    if p.startswith("Sample"):
+                        sample = int(p.replace("Sample", ""))
+                        break
+                else:
+                    logger.warning(
+                        f"Sample could not be identified for {transferred_file}"
+                    )
+                    return
+                capture_post(
+                    base_url=str(environment.url.geturl()),
+                    router_name="session_control.spa_router",
+                    function_name="register_atlas",
+                    token=self._token,
+                    instrument_name=environment.instrument_name,
+                    session_id=environment.murfey_session,
+                    data={
+                        "name": f"{environment.visit}-sample-{sample}",
+                        "acquisition_uuid": environment.acquisition_uuid,
+                        "register_grid": True,
+                        "tag": str(transferred_file.parent),
+                    },
+                )

src/murfey/client/contexts/spa_metadata.py

@@ -152,7 +152,7 @@ def post_transfer(
                                     "angle": pos_data[6],
                                 },
                             )
-                    if pos_data:
+                    if gs_pix_positions:
                         capture_post(
                             base_url=str(environment.url.geturl()),
                             router_name="session_control.spa_router",

src/murfey/server/api/session_control.py

@@ -359,7 +359,7 @@ def get_foil_hole(
 
 class AtlasRegistration(BaseModel):
     name: str
-    acquisition_uuid: str
+    acquisition_uuid: str | None
     storage_folder: str = ""
     register_grid: bool = False
     tag: str = ""
@@ -371,7 +371,7 @@ def register_atlas(
     atlas_registration_data: AtlasRegistration,
     db=murfey_db,
 ):
-    if SMARTEM_ACTIVE:
+    if SMARTEM_ACTIVE and atlas_registration_data.acquisition_uuid is not None:
         session = db.exec(select(Session).where(Session.id == session_id)).one()
         machine_config = get_machine_config(session.instrument_name)[
             session.instrument_name
@@ -382,7 +382,7 @@ def register_atlas(
             )
             grid_uuid = None
             if atlas_registration_data.tag:
-                dcg = murfey_db.exec(
+                dcg = db.exec(
                     select(DataCollectionGroup)
                     .where(DataCollectionGroup.session_id == session_id)
                     .where(DataCollectionGroup.tag == atlas_registration_data.tag)
@@ -410,7 +410,9 @@ def register_atlas(
                 if atlas_registration_data.register_grid:
                     smartem_client.grid_registered(grid_uuid)
     else:
-        logger.info("smartem deactivated so did not register atlas")
+        logger.info(
+            f"smartem deactivated so did not register atlas for {sanitise(str(atlas_registration_data.acquisition_uuid))}"
+        )
 
 
 class SquareRegistration(BaseModel):

src/murfey/server/api/workflow.py

@@ -232,6 +232,7 @@ def register_dc_group(
         )
     ).all():
         # Case where we switch from atlas to processing
+        original_tag = dcg_murfey[0].tag
         dcg_murfey[0].tag = dcg_params.tag or dcg_murfey[0].tag
         if _transport_object:
             _transport_object.send(
@@ -243,6 +244,13 @@ def register_dc_group(
                 },
             )
         db.add(dcg_murfey[0])
+        for grid_square in db.exec(
+            select(GridSquare)
+            .where(GridSquare.tag == original_tag)
+            .where(GridSquare.session_id == session_id)
+        ).all():
+            grid_square.tag = dcg_params.tag or original_tag
+            db.add(grid_square)
         db.commit()
     else:
         dcg_parameters = {

src/murfey/server/feedback.py

@@ -25,6 +25,7 @@
 from sqlalchemy import func
 from sqlalchemy.exc import (
     InvalidRequestError,
+    NoResultFound,
     OperationalError,
     PendingRollbackError,
     SQLAlchemyError,
@@ -353,7 +354,6 @@ def _release_2d_hold(message: dict, _db):
                 "nr_classes": default_spa_parameters.nr_classes_2d,
                 "do_icebreaker_jobs": default_spa_parameters.do_icebreaker_jobs,
                 "class2d_fraction_of_classes_to_remove": default_spa_parameters.fraction_of_classes_to_remove_2d,
-                "picker_id": feedback_params.picker_ispyb_id,
                 "class_uuids": _2d_class_murfey_ids(
                     first_class2d.particles_file, message["program_id"], _db
                 ),
@@ -431,7 +431,6 @@ def _release_3d_hold(message: dict, _db):
                 "mask_diameter": relion_params.mask_diameter or 0,
                 "do_initial_model": False if feedback_params.initial_model else True,
                 "initial_model_file": feedback_params.initial_model,
-                "picker_id": feedback_params.picker_ispyb_id,
                 "class_uuids": _3d_class_murfey_ids(
                     class3d_params.particles_file, _app_id(pj_id, _db), _db
                 ),
@@ -517,7 +516,6 @@ def _release_refine_hold(message: dict, _db):
                 "symmetry": relion_params.symmetry,
                 "node_creator_queue": machine_config.node_creator_queue,
                 "nr_iter": default_spa_parameters.nr_iter_3d,
-                "picker_id": feedback_params.picker_ispyb_id,
                 "refined_class_uuid": _refine_murfey_id(
                     refine_dir=refine_params.refine_dir,
                     tag=refine_params.tag,
@@ -589,13 +587,6 @@ def _register_incomplete_2d_batch(message: dict, _db):
     feedback_params.hold_class2d = True
     relion_options = dict(relion_params)
     other_options = dict(feedback_params)
-    if other_options["picker_ispyb_id"] is None:
-        logger.info("No ISPyB particle picker ID yet")
-        feedback_params.hold_class2d = False
-        _db.add(feedback_params)
-        _db.commit()
-        _db.expunge(feedback_params)
-        return
     _db.add(feedback_params)
     _db.commit()
     _db.expunge(feedback_params)
@@ -629,7 +620,6 @@ def _register_incomplete_2d_batch(message: dict, _db):
             "batch_is_complete": False,
             "particle_diameter": relion_options["particle_diameter"],
             "combine_star_job_number": -1,
-            "picker_id": other_options["picker_ispyb_id"],
             "nr_iter": default_spa_parameters.nr_iter_2d,
             "batch_size": default_spa_parameters.batch_size_2d,
             "nr_classes": default_spa_parameters.nr_classes_2d,
@@ -693,7 +683,7 @@ def _register_complete_2d_batch(message: dict, _db):
     ).one()
     _db.expunge(relion_params)
     _db.expunge(feedback_params)
-    if feedback_params.hold_class2d or feedback_params.picker_ispyb_id is None:
+    if feedback_params.hold_class2d:
         feedback_params.rerun_class2d = True
         _db.add(feedback_params)
         _db.commit()
@@ -791,7 +781,6 @@ def _register_complete_2d_batch(message: dict, _db):
                 "mask_diameter": relion_params.mask_diameter or 0,
                 "combine_star_job_number": feedback_params.star_combination_job,
                 "autoselect_min_score": 0,
-                "picker_id": feedback_params.picker_ispyb_id,
                 "class_uuids": class_uuids,
                 "class2d_grp_uuid": class2d_grp_uuid,
                 "nr_iter": default_spa_parameters.nr_iter_2d,
@@ -866,7 +855,6 @@ def _register_complete_2d_batch(message: dict, _db):
                 "mask_diameter": relion_params.mask_diameter or 0,
                 "combine_star_job_number": feedback_params.star_combination_job,
                 "autoselect_min_score": feedback_params.class_selection_score or 0,
-                "picker_id": feedback_params.picker_ispyb_id,
                 "class_uuids": class_uuids,
                 "class2d_grp_uuid": class2d_grp_uuid,
                 "nr_iter": default_spa_parameters.nr_iter_2d,
@@ -957,7 +945,6 @@ def _flush_class2d(
                 "mask_diameter": relion_params.mask_diameter or 0,
                 "combine_star_job_number": feedback_params.star_combination_job,
                 "autoselect_min_score": feedback_params.class_selection_score or 0,
-                "picker_id": feedback_params.picker_ispyb_id,
                 "class_uuids": _2d_class_murfey_ids(
                     saved_message.particles_file, _app_id(pj_id, _db), _db
                 ),
@@ -1007,16 +994,6 @@ def _register_class_selection(message: dict, _db):
     ).one()
     _db.expunge(feedback_params)
 
-    if feedback_params.picker_ispyb_id is None:
-        selection_stash = db.SelectionStash(
-            pj_id=pj_id,
-            class_selection_score=message["class_selection_score"] or 0,
-        )
-        _db.add(selection_stash)
-        _db.commit()
-        _db.close()
-        return
-
     feedback_params.class_selection_score = message.get("class_selection_score") or 0
     feedback_params.hold_class2d = False
     _flush_class2d(
@@ -1284,7 +1261,6 @@ def _register_3d_batch(message: dict, _db):
                 "particle_diameter": relion_options["particle_diameter"],
                 "mask_diameter": relion_options["mask_diameter"] or 0,
                 "do_initial_model": True,
-                "picker_id": other_options["picker_ispyb_id"],
                 "class_uuids": {i + 1: m for i, m in enumerate(class_uuids)},
                 "class3d_grp_uuid": class3d_grp_uuid,
                 "nr_iter": default_spa_parameters.nr_iter_3d,
@@ -1325,7 +1301,6 @@ def _register_3d_batch(message: dict, _db):
                 "mask_diameter": relion_options["mask_diameter"] or 0,
                 "do_initial_model": False,
                 "initial_model_file": other_options["initial_model"],
-                "picker_id": other_options["picker_ispyb_id"],
                 "class_uuids": _3d_class_murfey_ids(
                     class3d_params.particles_file, _app_id(pj_id, _db), _db
                 ),
@@ -1534,7 +1509,6 @@ def _register_refinement(message: dict, _db):
             db.ClassificationFeedbackParameters.pj_id == pj_id_params
         )
     ).one()
-    other_options = dict(feedback_params)
 
     if feedback_params.hold_refine:
         # If waiting then save the message
@@ -1626,7 +1600,6 @@ def _register_refinement(message: dict, _db):
                 "symmetry": relion_options["symmetry"],
                 "node_creator_queue": machine_config.node_creator_queue,
                 "nr_iter": default_spa_parameters.nr_iter_3d,
-                "picker_id": other_options["picker_ispyb_id"],
                 "refined_class_uuid": _refine_murfey_id(
                     refine_dir=refine_params.refine_dir,
                     tag=refine_params.tag,
@@ -1761,7 +1734,6 @@ def _register_bfactors(message: dict, _db):
                 "particle_diameter": relion_options["particle_diameter"],
                 "mask_diameter": relion_options["mask_diameter"] or 0,
                 "node_creator_queue": machine_config.node_creator_queue,
-                "picker_id": feedback_params.picker_ispyb_id,
                 "refined_grp_uuid": bfactor_params.refined_grp_uuid,
                 "refined_class_uuid": bfactor_params.refined_class_uuid,
                 "session_id": message["session_id"],
@@ -2189,6 +2161,12 @@ def feedback_callback(header: dict, message: dict, _db=murfey_db) -> None:
         time.sleep(1)
         if murfey.server._transport_object:
             murfey.server._transport_object.transport.nack(header, requeue=True)
+    except NoResultFound:
+        # Missing rows might be due to a race condition and should be requeued
+        logger.warning("No matching database row was found", exc_info=True)
+        time.sleep(1)
+        if murfey.server._transport_object:
+            murfey.server._transport_object.transport.nack(header, requeue=True)
     except Exception:
         logger.warning(
             "Exception encountered in server RabbitMQ callback", exc_info=True

src/murfey/util/db.py

@@ -339,9 +339,6 @@ class ProcessingJob(SQLModel, table=True):  # type: ignore
     auto_proc_programs: List["AutoProcProgram"] = Relationship(
         back_populates="processing_job", sa_relationship_kwargs={"cascade": "delete"}
     )
-    selection_stash: List["SelectionStash"] = Relationship(
-        back_populates="processing_job", sa_relationship_kwargs={"cascade": "delete"}
-    )
     particle_sizes: List["ParticleSizes"] = Relationship(
         back_populates="processing_job", sa_relationship_kwargs={"cascade": "delete"}
     )
@@ -393,15 +390,6 @@ class PreprocessStash(SQLModel, table=True):  # type: ignore
     foil_hole: Optional["FoilHole"] = Relationship(back_populates="preprocess_stashes")
 
 
-class SelectionStash(SQLModel, table=True):  # type: ignore
-    id: Optional[int] = Field(default=None, primary_key=True)
-    class_selection_score: float
-    pj_id: int = Field(foreign_key="processingjob.processingJobId")
-    processing_job: Optional[ProcessingJob] = Relationship(
-        back_populates="selection_stash"
-    )
-
-
 class TomographyProcessingParameters(SQLModel, table=True):  # type: ignore
     dcg_id: int = Field(
         primary_key=True, foreign_key="datacollectiongroup.dataCollectionGroupId"
@@ -717,7 +705,6 @@ class ClassificationFeedbackParameters(SQLModel, table=True):  # type: ignore
     initial_model: str
     next_job: int
     picker_murfey_id: Optional[int] = Field(default=None, foreign_key="murfeyledger.id")
-    picker_ispyb_id: Optional[int] = None
     processing_job: Optional[ProcessingJob] = Relationship(
         back_populates="classification_feedback_parameters"
     )

src/murfey/util/processing_params.py

@@ -1,7 +1,6 @@
 from datetime import datetime
 from functools import lru_cache
 from pathlib import Path
-from typing import Literal, Optional
 
 from pydantic import BaseModel
 from werkzeug.utils import secure_filename
@@ -60,12 +59,6 @@ class CLEMProcessingParameters(BaseModel):
     # Atlas vs GridSquare registration threshold
     atlas_threshold: float = 0.0015  # in m
 
-    # Image alignment and merging-specific parameters
-    crop_to_n_frames: Optional[int] = 50
-    align_self: Literal["enabled", ""] = "enabled"
-    flatten: Literal["mean", "min", "max", ""] = "mean"
-    align_across: Literal["enabled", ""] = "enabled"
-
 
 default_clem_processing_parameters = CLEMProcessingParameters()
 

src/murfey/workflows/clem/align_and_merge.py

@@ -6,7 +6,7 @@
 from __future__ import annotations
 
 from pathlib import Path
-from typing import Literal, Optional
+from typing import Optional
 
 from murfey.util.config import get_machine_config
 
@@ -24,11 +24,6 @@ def run(
     series_name: str,
     images: list[Path],
     metadata: Path,
-    # Optional processing parameters
-    crop_to_n_frames: Optional[int] = None,
-    align_self: Literal["enabled", ""] = "",
-    flatten: Literal["mean", "min", "max", ""] = "mean",
-    align_across: Literal["enabled", ""] = "",
     # Optional session parameters
     messenger: Optional[TransportManager] = None,
 ):
@@ -65,10 +60,6 @@ def run(
                 "series_name": series_name,
                 "images": [str(file) for file in images],
                 "metadata": str(metadata),
-                "crop_to_n_frames": crop_to_n_frames,
-                "align_self": align_self,
-                "flatten": flatten,
-                "align_across": align_across,
                 # Other recipe parameters
                 "session_dir": str(session_dir),
                 "session_id": session_id,