When registering a new data collection group check other data collections in the group for matching grid square IDs and transfer grid squares to new data collection group

Both data collection groups will use the same data collection group IDs but will differ by tag

Author: d-j-hatton

src/murfey/client/contexts/spa_metadata.py

@@ -176,6 +176,9 @@ def post_transfer(
                     logger.warning(f"Cannot find Images-Disc* in {dcg_search_dir}")
                     return
                 dcg_tag = str(dcg_images_dirs[-1])
+                gs_pix_positions = get_grid_square_atlas_positions(
+                    source_visit_dir / partial_path
+                )
                 dcg_data = {
                     "experiment_type": "single particle",
                     "experiment_type_id": 37,
@@ -187,11 +190,11 @@ def post_transfer(
                     ),
                     "sample": environment.samples[source].sample,
                     "atlas_pixel_size": atlas_pixel_size,
+                    "grid_squares_to_match": [
+                        int(gs) for gs in gs_pix_positions.keys()
+                    ],
                 }
                 capture_post(url, json=dcg_data)
-                gs_pix_positions = get_grid_square_atlas_positions(
-                    source_visit_dir / partial_path
-                )
                 for gs, pos_data in gs_pix_positions.items():
                     if pos_data:
                         capture_post(
@@ -228,14 +231,15 @@ def post_transfer(
                 logger.warning(f"Cannot find Images-Disc* in {dcg_search_dir}")
                 return
             dcg_tag = str(dcg_images_dirs[-1])
+            gs_name = transferred_file.stem.split("_")[1]
             dcg_data = {
                 "experiment_type": "single particle",
                 "experiment_type_id": 37,
                 "tag": dcg_tag,
+                "grid_squares_to_match": [int(gs_name)],
             }
             capture_post(url, json=dcg_data)
 
-            gs_name = transferred_file.stem.split("_")[1]
             logger.info(
                 f"Collecting foil hole positions for {str(transferred_file)} and grid square {int(gs_name)}"
             )

src/murfey/server/__init__.py

@@ -2508,6 +2508,8 @@ def feedback_callback(header: dict, message: dict) -> None:
                 .where(db.DataCollectionGroup.tag == message.get("tag"))
             ).all():
                 dcgid = dcg_murfey[0].id
+            elif message.get("data_collection_group_id") is not None:
+                dcgid = message["data_collection_group_id"]
             else:
                 if ispyb_session_id is None:
                     murfey_dcg = db.DataCollectionGroup(
@@ -2575,8 +2577,8 @@ def feedback_callback(header: dict, message: dict) -> None:
                     message["atlas"],
                     message["atlas_pixel_size"],
                     message["sample"],
-                    experiment_type=message["experiment_type"],
-                    experiment_type_id=message["experiment_type_id"],
+                    experiment_type=message.get("experiment_type", ""),
+                    experiment_type_id=message.get("experiment_type_id"),
                 )
                 _transport_object.transport.ack(header)
             return None

src/murfey/server/api/__init__.py

@@ -51,6 +51,7 @@
 from murfey.server.gain import Camera, prepare_eer_gain, prepare_gain
 from murfey.server.murfey_db import murfey_db
 from murfey.util import secure_path
+from murfey.util.atlas import check_for_common_grid_squares, replace_grid_square_tags
 from murfey.util.config import MachineConfig, from_file, settings
 from murfey.util.db import (
     AutoProcProgram,
@@ -1301,6 +1302,8 @@ def get_data_collections(
 def register_dc_group(
     visit_name, session_id: MurfeySessionID, dcg_params: DCGroupParameters, db=murfey_db
 ):
+    # need to check for matching grid squares if nothing is found based on tag
+    # if that matches then grid square tags need reassigning
     ispyb_proposal_code = visit_name[:2]
     ispyb_proposal_number = visit_name.split("-")[0][2:]
     ispyb_visit_number = visit_name.split("-")[-1]
@@ -1343,6 +1346,28 @@ def register_dc_group(
                 dcg_murfey[0].atlas_id = atlas_id_response["return_value"]
         db.add(dcg_murfey[0])
         db.commit()
+    elif matching_dcg := check_for_common_grid_squares(
+        murfey_db,
+        dcg_params.tag,
+        session_id,
+        grid_square_names=dcg_params.grid_squares_to_match,
+    ):
+        dcg_parameters = {
+            "data_collection_group_id": matching_dcg[0],
+            "start_time": str(datetime.datetime.now()),
+            "experiment_type": dcg_params.experiment_type,
+            "experiment_type_id": dcg_params.experiment_type_id,
+            "tag": dcg_params.tag,
+            "session_id": session_id,
+            "atlas": dcg_params.atlas,
+            "sample": dcg_params.sample,
+            "atlas_pixel_size": dcg_params.atlas_pixel_size,
+        }
+        if _transport_object:
+            _transport_object.send(
+                _transport_object.feedback_queue, {"register": "data_collection_group", **dcg_parameters, "microscope": instrument_name, "proposal_code": ispyb_proposal_code, "proposal_number": ispyb_proposal_number, "visit_number": ispyb_visit_number}  # type: ignore
+            )
+        replace_grid_square_tags(murfey_db, matching_dcg[1], dcg_params.tag)
     else:
         dcg_parameters = {
             "start_time": str(datetime.datetime.now()),

src/murfey/util/atlas.py

@@ -0,0 +1,46 @@
+from typing import List, Optional, Tuple
+
+from sqlmodel import Session, select
+
+from murfey.util.db import DataCollectionGroup, GridSquare
+
+
+def check_for_common_grid_squares(
+    murfey_db: Session,
+    tag: str,
+    session_id: int,
+    grid_square_names: Optional[List[int]] = None,
+) -> Tuple[Optional[int], str]:
+    data_collection_groups = murfey_db.exec(
+        select(DataCollectionGroup).where(DataCollectionGroup.session_id == session_id)
+    ).all()
+    grid_squares_matching_tag = [
+        gs.name
+        for gs in murfey_db.exec(
+            select(GridSquare)
+            .where(GridSquare.tag == tag)
+            .where(GridSquare.session_id == session_id)
+        ).all()
+    ]
+    if grid_square_names is not None:
+        grid_squares_matching_tag.extend(grid_square_names)
+    for dcg in data_collection_groups:
+        grid_squares = murfey_db.exec(
+            select(GridSquare).where(GridSquare.tag == dcg.tag)
+        ).all()
+        if set(grid_squares_matching_tag).intersection(
+            {gs.name for gs in grid_squares}
+        ):
+            return (dcg.id, dcg.tag)
+    return (None, "")
+
+
+def replace_grid_square_tags(murfey_db: Session, old_tag: str, new_tag: str):
+    grid_squares = murfey_db.exec(
+        select(GridSquare).where(GridSquare.tag == old_tag)
+    ).all()
+    for gs in grid_squares:
+        gs.tag = new_tag
+    murfey_db.add(grid_squares)
+    murfey_db.commit()
+    murfey_db.close()

src/murfey/util/db.py

@@ -348,7 +348,7 @@ class Tilt(SQLModel, table=True):  # type: ignore
 
 
 class DataCollectionGroup(SQLModel, table=True):  # type: ignore
-    id: int = Field(primary_key=True, unique=True)
+    id: int = Field(primary_key=True)
     session_id: int = Field(foreign_key="session.id", primary_key=True)
     tag: str = Field(primary_key=True)
     atlas_id: Optional[int] = None

src/murfey/util/models.py

@@ -48,6 +48,7 @@ class DCGroupParameters(BaseModel):
     atlas: str = ""
     sample: Optional[int] = None
     atlas_pixel_size: int = 0
+    grid_squares_to_match: List[int] = []
 
 
 class DCParameters(BaseModel):