Match grid squares on sample rather than tag if sample known

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