Match grid squares on sample rather than tag if sample known (#827)

There is a problem with registering grid squares from an atlas, as the dcg tag is being used as the matching but this can change if atlas converts to SPA.

Instead, when registering squares from an atlas file, provide the sample and match this to a dcg. And send dcg request before that to make sure one will exist.

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 is not None:
+        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(

tests/client/contexts/test_atlas.py

@@ -139,7 +139,21 @@ def test_atlas_context_dm(mock_capture_post, tmp_path):
     context = AtlasContext("tomo", tmp_path, {}, "token")
     context.post_transfer(atlas_dm, environment=env)
 
-    assert mock_capture_post.call_count == 5
+    assert mock_capture_post.call_count == 6
+    mock_capture_post.assert_any_call(
+        base_url="http://localhost:8000",
+        router_name="workflow.router",
+        function_name="register_dc_group",
+        token="token",
+        instrument_name="m01",
+        visit_name="cm12345-6",
+        session_id=1,
+        data={
+            "experiment_type_id": 44,  # Atlas
+            "tag": str(atlas_dm.parent),
+            "sample": 2,
+        },
+    )
     mock_capture_post.assert_any_call(
         base_url="http://localhost:8000",
         router_name="session_control.spa_router",
@@ -150,6 +164,7 @@ def test_atlas_context_dm(mock_capture_post, tmp_path):
         gsid=101,
         data={
             "tag": str(atlas_dm.parent),
+            "sample": 2,
             "x_location": 1200,
             "y_location": 1500,
             "x_stage_position": 2e9,

tests/workflows/spa/test_flush_spa_preprocess.py

@@ -14,6 +14,14 @@ def test_register_grid_square_update_add_locations(
 ):
     """Test the updating of an existing grid square"""
     # Create a grid square to update
+    dcg = DataCollectionGroup(
+        id=1,
+        session_id=ExampleVisit.murfey_session_id,
+        tag="session_tag",
+        atlas_id=90,
+        sample=2,
+    )
+    murfey_db_session.add(dcg)
     grid_square = GridSquare(
         id=1,
         name=101,
@@ -58,6 +66,14 @@ def test_register_grid_square_update_add_nothing(
 ):
     """Test the updating of an existing grid square, but with nothing to update with"""
     # Create a grid square to update
+    dcg = DataCollectionGroup(
+        id=1,
+        session_id=ExampleVisit.murfey_session_id,
+        tag="session_tag",
+        atlas_id=90,
+        sample=2,
+    )
+    murfey_db_session.add(dcg)
     grid_square = GridSquare(
         id=1,
         name=101,
@@ -71,8 +87,8 @@ def test_register_grid_square_update_add_nothing(
     murfey_db_session.add(grid_square)
     murfey_db_session.commit()
 
-    # Parameters to update with
-    new_parameters = GridSquareParameters(tag="session_tag")
+    # Parameters to update with - use a new tag but same sample
+    new_parameters = GridSquareParameters(tag="atlas_tag", sample=2)
 
     # Run the registration
     flush_spa_preprocess.register_grid_square(
@@ -95,13 +111,14 @@ def test_register_grid_square_insert_with_ispyb(
     mock_transport, murfey_db_session: Session, tmp_path
 ):
     # Create a data collection group for lookups
-    grid_square = DataCollectionGroup(
+    dcg = DataCollectionGroup(
         id=1,
         session_id=ExampleVisit.murfey_session_id,
         tag="session_tag",
         atlas_id=90,
+        sample=2,
     )
-    murfey_db_session.add(grid_square)
+    murfey_db_session.add(dcg)
     murfey_db_session.commit()
 
     # Set the ispyb return