Pass sample to atlas square registration for safety when tag changes

Author: stephen-riggs

src/murfey/client/contexts/atlas.py

@@ -167,6 +167,30 @@ def post_transfer_epu(
             except KeyError:
                 logger.info("Unable to read grid square locations from Atlas.dm")
                 return
+            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
+
+            # Make sure a dcg is requested before doing grid squares
+            capture_post(
+                base_url=str(environment.url.geturl()),
+                router_name="workflow.router",
+                function_name="register_dc_group",
+                token=self._token,
+                instrument_name=environment.instrument_name,
+                visit_name=environment.visit,
+                session_id=environment.murfey_session,
+                data={
+                    "experiment_type_id": 44,  # Atlas
+                    "tag": str(transferred_file.parent),
+                    "sample": sample,
+                },
+            )
+            # Register all grid squares on this atlas
             for gs, pos_data in gs_pix_positions.items():
                 if pos_data:
                     capture_post(
@@ -179,6 +203,7 @@ def post_transfer_epu(
                         gsid=int(gs),
                         data={
                             "tag": str(transferred_file.parent),
+                            "sample": sample,
                             "x_location": pos_data[0],
                             "y_location": pos_data[1],
                             "x_stage_position": pos_data[2],
@@ -188,16 +213,8 @@ def post_transfer_epu(
                             "angle": pos_data[6],
                         },
                     )
+            # Register atlas in smartem
             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",

src/murfey/util/models.py

@@ -272,6 +272,7 @@ class Base(BaseModel):
 class GridSquareParameters(BaseModel):
     tag: str
     image: str = ""
+    sample: int | None = None
 
     # Actual coordinates for image centre in real space
     x_location: Optional[float] = None

src/murfey/workflows/spa/flush_spa_preprocess.py

@@ -65,10 +65,22 @@ def register_grid_square(
     if grid_square_params.width is not None:
         grid_square_params.width_scaled = int(grid_square_params.width / 7.8)
 
+    if grid_square_params.sample:
+        dcg = murfey_db.exec(
+            select(DataCollectionGroup)
+            .where(DataCollectionGroup.session_id == session_id)
+            .where(DataCollectionGroup.sample == grid_square_params.sample)
+        ).one()
+    else:
+        dcg = murfey_db.exec(
+            select(DataCollectionGroup)
+            .where(DataCollectionGroup.session_id == session_id)
+            .where(DataCollectionGroup.tag == grid_square_params.tag)
+        ).one()
     grid_square_query = murfey_db.exec(
         select(GridSquare)
         .where(GridSquare.name == gsid)
-        .where(GridSquare.tag == grid_square_params.tag)
+        .where(GridSquare.tag == dcg.tag)
         .where(GridSquare.session_id == session_id)
     ).all()
     if grid_square_query:
@@ -101,11 +113,6 @@ def register_grid_square(
     else:
         # No existing grid square in the murfey database
         if _transport_object:
-            dcg = murfey_db.exec(
-                select(DataCollectionGroup)
-                .where(DataCollectionGroup.session_id == session_id)
-                .where(DataCollectionGroup.tag == grid_square_params.tag)
-            ).one()
             gs_ispyb_response = _transport_object.do_insert_grid_square(
                 dcg.atlas_id, gsid, grid_square_params
             )
@@ -149,12 +156,7 @@ def register_grid_square(
                 instrument_name=murfey_session.instrument_name
             )[murfey_session.instrument_name]
             if machine_config.smartem_api_url:
-                dcg = murfey_db.exec(
-                    select(DataCollectionGroup)
-                    .where(DataCollectionGroup.session_id == session_id)
-                    .where(DataCollectionGroup.tag == grid_square_params.tag)
-                ).one_or_none()
-                if dcg and dcg.smartem_grid_uuid:
+                if dcg.smartem_grid_uuid:
                     secured_grid_square_image_path_full_res: Path | None = None
                     if grid_square_params.image:
                         secured_grid_square_image_path_full_res = secure_path(