Register all grid squares from atlases (#780)

Currently we only register EPU grid squares when data collection is started. This sets up a method to register them in all atlases as during screening.

The problem comes on conversion of screening sessions to EPU sessions. Here this loops through all grid squares and changes the tag when the data collection experiment type id is changed.

Author: stephen-riggs

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/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)

tests/client/contexts/test_atlas.py

@@ -33,10 +33,7 @@ def test_atlas_context_mrc(mock_capture_post, tmp_path):
     atlas_mrc.parent.mkdir(parents=True)
     atlas_mrc.touch()
 
-    context.post_transfer(
-        atlas_mrc,
-        environment=env,
-    )
+    context.post_transfer(atlas_mrc, environment=env)
     mock_capture_post.assert_called_once_with(
         base_url="http://localhost:8000",
         router_name="session_control.spa_router",
@@ -72,10 +69,7 @@ def test_atlas_context_xml(mock_capture_post, tmp_path):
             "</numericValue></x></pixelSize></SpatialScale></MicroscopeImage>"
         )
 
-    context.post_transfer(
-        atlas_xml,
-        environment=env,
-    )
+    context.post_transfer(atlas_xml, environment=env)
     dcg_data = {
         "experiment_type_id": 44,  # Atlas
         "tag": str(atlas_xml.parent),
@@ -95,3 +89,87 @@ def test_atlas_context_xml(mock_capture_post, tmp_path):
         session_id=1,
         data=dcg_data,
     )
+
+
+@patch("murfey.client.contexts.atlas.capture_post")
+def test_atlas_context_dm(mock_capture_post, tmp_path):
+    env = MurfeyInstanceEnvironment(
+        url=urlparse("http://localhost:8000"),
+        client_id=0,
+        sources=[tmp_path / "cm12345-6"],
+        default_destinations={
+            tmp_path / "cm12345-6": f"{tmp_path}/destination/cm12345-6"
+        },
+        instrument_name="m01",
+        visit="cm12345-6",
+        murfey_session=1,
+        acquisition_uuid="uuid1",
+    )
+
+    # Write sample dm file
+    atlas_dm = tmp_path / "cm12345-6/Supervisor_atlas/Sample2/Atlas/Atlas.dm"
+    atlas_dm.parent.mkdir(parents=True)
+    grid_square_values = (
+        "<value><b:PositionOnTheAtlas>"
+        "<c:Center><d:x>1200</d:x><d:y>1500</d:y></c:Center>"
+        "<c:Physical><d:x>2</d:x><d:y>3</d:y></c:Physical>"
+        "<c:Size><d:width>130</d:width><d:height>560</d:height></c:Size>"
+        "<c:Rotation>0.14</c:Rotation>"
+        "</b:PositionOnTheAtlas></value>"
+    )
+    with open(atlas_dm, "w") as new_xml:
+        new_xml.write(
+            "<AtlasSessionXml><Atlas><TilesEfficient><_items>"
+            # First tile with two grid squares
+            "<TileXml><Nodes><KeyValuePairs><KeyValuePairOfintNodeXml1>"
+            f"<key>101</key>{grid_square_values}"
+            "</KeyValuePairOfintNodeXml1><KeyValuePairOfintNodeXml1>"
+            f"<key>102</key>{grid_square_values}"
+            "</KeyValuePairOfintNodeXml1></KeyValuePairs></Nodes></TileXml>"
+            # Second tile with two grid squares
+            "<TileXml><Nodes><KeyValuePairs><KeyValuePairOfintNodeXml1>"
+            f"<key>103</key>{grid_square_values}"
+            "</KeyValuePairOfintNodeXml1><KeyValuePairOfintNodeXml1>"
+            f"<key>104</key>{grid_square_values}"
+            "</KeyValuePairOfintNodeXml1></KeyValuePairs></Nodes></TileXml>"
+            # Close all
+            "</_items></TilesEfficient></Atlas></AtlasSessionXml>"
+        )
+
+    context = AtlasContext("tomo", tmp_path, {}, "token")
+    context.post_transfer(atlas_dm, environment=env)
+
+    assert mock_capture_post.call_count == 5
+    mock_capture_post.assert_any_call(
+        base_url="http://localhost:8000",
+        router_name="session_control.spa_router",
+        function_name="register_grid_square",
+        token="token",
+        instrument_name="m01",
+        session_id=1,
+        gsid=101,
+        data={
+            "tag": str(atlas_dm.parent),
+            "x_location": 1200,
+            "y_location": 1500,
+            "x_stage_position": 2e9,
+            "y_stage_position": 3e9,
+            "width": 130,
+            "height": 560,
+            "angle": 0.14,
+        },
+    )
+    mock_capture_post.assert_any_call(
+        base_url="http://localhost:8000",
+        router_name="session_control.spa_router",
+        function_name="register_atlas",
+        token="token",
+        instrument_name="m01",
+        session_id=1,
+        data={
+            "name": "cm12345-6-sample-2",
+            "acquisition_uuid": "uuid1",
+            "register_grid": True,
+            "tag": str(atlas_dm.parent),
+        },
+    )