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

@@ -354,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
                 ),
@@ -432,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
                 ),
@@ -518,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,
@@ -590,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)
@@ -630,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,
@@ -694,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()
@@ -792,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,
@@ -867,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,
@@ -958,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
                 ),
@@ -1008,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(
@@ -1285,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,
@@ -1326,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
                 ),
@@ -1535,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
@@ -1627,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,
@@ -1762,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"],

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/workflows/spa/flush_spa_preprocess.py

@@ -297,6 +297,7 @@ def register_foil_hole(
                 )
                 fh_data = SmartEMFoilHoleData(
                     id=str(foil_hole_params.name),
+                    gridsquare_id=str(gs.name),
                     gridsquare_uuid=gs.smartem_uuid,
                     x_location=(
                         int(foil_hole_params.x_location)