Match data collection groups for atlases and processing (#696)

Creates a common function for data collection group creation by the client. The data collection groups are made as atlases then updated to be SPA or Tomo if data collection happens.

Diff looks bigger than it really is as a test file got moved.

Author: stephen-riggs

pyproject.toml

@@ -108,6 +108,7 @@ GitHub = "https://github.com/DiamondLightSource/python-murfey"
 "clem.register_preprocessing_result" = "murfey.workflows.clem.register_preprocessing_results:run"
 "data_collection" = "murfey.workflows.register_data_collection:run"
 "data_collection_group" = "murfey.workflows.register_data_collection_group:run"
+"experiment_type_update" = "murfey.workflows.register_experiment_type_update:run"
 "pato" = "murfey.workflows.notifications:notification_setup"
 "picked_particles" = "murfey.workflows.spa.picking:particles_picked"
 "picked_tomogram" = "murfey.workflows.tomo.picking:picked_tomogram"

src/murfey/client/context.py

@@ -3,16 +3,156 @@
 import logging
 from importlib.metadata import entry_points
 from pathlib import Path
-from typing import Any, Dict, List, NamedTuple
+from typing import Any, List, NamedTuple
 
-from murfey.client.instance_environment import MurfeyInstanceEnvironment
+import xmltodict
+
+from murfey.client.instance_environment import MurfeyInstanceEnvironment, SampleInfo
+from murfey.util.client import capture_post, get_machine_config_client
 
 logger = logging.getLogger("murfey.client.context")
 
 
-class FutureRequest(NamedTuple):
-    url: str
-    message: Dict[str, Any]
+def _atlas_destination(
+    environment: MurfeyInstanceEnvironment, source: Path, token: str
+) -> Path:
+    machine_config = get_machine_config_client(
+        str(environment.url.geturl()),
+        token,
+        instrument_name=environment.instrument_name,
+        demo=environment.demo,
+    )
+    for i, destination_part in enumerate(
+        Path(environment.default_destinations[source]).parts
+    ):
+        if destination_part == environment.visit:
+            return Path(machine_config.get("rsync_basepath", "")) / "/".join(
+                Path(environment.default_destinations[source]).parent.parts[: i + 1]
+            )
+    return (
+        Path(machine_config.get("rsync_basepath", ""))
+        / Path(environment.default_destinations[source]).parent
+        / environment.visit
+    )
+
+
+def ensure_dcg_exists(
+    collection_type: str,
+    metadata_source: Path,
+    environment: MurfeyInstanceEnvironment,
+    token: str,
+) -> str | None:
+    """Create  a data collection group"""
+    if collection_type == "tomo":
+        experiment_type_id = 36
+        session_file = metadata_source / "Session.dm"
+    elif collection_type == "spa":
+        experiment_type_id = 37
+        session_file = metadata_source / "EpuSession.dm"
+        for h in entry_points(group="murfey.hooks"):
+            try:
+                if h.name == "get_epu_session_metadata":
+                    h.load()(session_file, environment=environment)
+            except Exception as e:
+                logger.warning(f"Get EPU session hook failed: {e}")
+    else:
+        logger.error(f"Unknown collection type {collection_type}")
+        return None
+
+    if not session_file.is_file():
+        logger.warning(f"Cannot find session file {str(session_file)}")
+        dcg_tag = (
+            str(metadata_source).replace(f"/{environment.visit}", "").replace("//", "/")
+        )
+        dcg_data = {
+            "experiment_type_id": experiment_type_id,
+            "tag": dcg_tag,
+        }
+    else:
+        with open(session_file, "r") as session_xml:
+            session_data = xmltodict.parse(session_xml.read())
+
+        if collection_type == "tomo":
+            windows_path = session_data["TomographySession"]["AtlasId"]
+        else:
+            windows_path = session_data["EpuSessionXml"]["Samples"]["_items"][
+                "SampleXml"
+            ][0]["AtlasId"]["#text"]
+
+        logger.info(f"Windows path to atlas metadata found: {windows_path}")
+        if not windows_path:
+            logger.warning("No atlas metadata path found")
+            return None
+        visit_index = windows_path.split("\\").index(environment.visit)
+        partial_path = "/".join(windows_path.split("\\")[visit_index + 1 :])
+        logger.info("Partial Linux path successfully constructed from Windows path")
+
+        source_visit_dir = metadata_source.parent
+        logger.info(
+            f"Looking for atlas XML file in metadata directory {str((source_visit_dir / partial_path).parent)}"
+        )
+        atlas_xml_path = list(
+            (source_visit_dir / partial_path).parent.glob("Atlas_*.xml")
+        )[0]
+        logger.info(f"Atlas XML path {str(atlas_xml_path)} found")
+        with open(atlas_xml_path, "rb") as atlas_xml:
+            atlas_xml_data = xmltodict.parse(atlas_xml)
+            atlas_original_pixel_size = float(
+                atlas_xml_data["MicroscopeImage"]["SpatialScale"]["pixelSize"]["x"][
+                    "numericValue"
+                ]
+            )
+        # need to calculate the pixel size of the downscaled image
+        atlas_pixel_size = atlas_original_pixel_size * 7.8
+        logger.info(f"Atlas image pixel size determined to be {atlas_pixel_size}")
+
+        for p in partial_path.split("/"):
+            if p.startswith("Sample"):
+                sample = int(p.replace("Sample", ""))
+                break
+        else:
+            logger.warning(f"Sample could not be identified for {metadata_source}")
+            return None
+        environment.samples[metadata_source] = SampleInfo(
+            atlas=Path(partial_path), sample=sample
+        )
+
+        dcg_search_dir = (
+            str(metadata_source).replace(f"/{environment.visit}", "").replace("//", "/")
+        )
+        if collection_type == "tomo":
+            dcg_tag = dcg_search_dir
+        else:
+            dcg_images_dirs = sorted(
+                Path(dcg_search_dir).glob("Images-Disc*"),
+                key=lambda x: x.stat().st_ctime,
+            )
+            if not dcg_images_dirs:
+                logger.warning(f"Cannot find Images-Disc* in {dcg_search_dir}")
+                return None
+            dcg_tag = str(dcg_images_dirs[-1])
+
+        dcg_data = {
+            "experiment_type_id": experiment_type_id,
+            "tag": dcg_tag,
+            "atlas": str(
+                _atlas_destination(environment, metadata_source, token)
+                / environment.samples[metadata_source].atlas.parent
+                / atlas_xml_path.with_suffix(".jpg").name
+            ).replace("//", "/"),
+            "sample": environment.samples[metadata_source].sample,
+            "atlas_pixel_size": atlas_pixel_size,
+        }
+    capture_post(
+        base_url=str(environment.url.geturl()),
+        router_name="workflow.router",
+        function_name="register_dc_group",
+        token=token,
+        visit_name=environment.visit,
+        session_id=environment.murfey_session,
+        data=dcg_data,
+    )
+    return dcg_tag
 
 
 class ProcessingParameter(NamedTuple):

src/murfey/client/contexts/atlas.py

@@ -4,9 +4,8 @@
 
 import xmltodict
 
-from murfey.client.context import Context
+from murfey.client.context import Context, _atlas_destination
 from murfey.client.contexts.spa import _get_source
-from murfey.client.contexts.spa_metadata import _atlas_destination
 from murfey.client.instance_environment import MurfeyInstanceEnvironment
 from murfey.util.client import capture_post
 
@@ -38,7 +37,7 @@ def post_transfer(
             source = _get_source(transferred_file, environment)
             if source:
                 transferred_atlas_name = _atlas_destination(
-                    environment, source, transferred_file, self._token
+                    environment, source, self._token
                 ) / transferred_file.relative_to(source.parent)
                 capture_post(
                     base_url=str(environment.url.geturl()),
@@ -60,7 +59,7 @@ def post_transfer(
             if source:
                 atlas_mrc = transferred_file.with_suffix(".mrc")
                 transferred_atlas_jpg = _atlas_destination(
-                    environment, source, atlas_mrc, self._token
+                    environment, source, self._token
                 ) / atlas_mrc.relative_to(source.parent).with_suffix(".jpg")
 
                 with open(transferred_file, "rb") as atlas_xml:

src/murfey/client/contexts/spa_metadata.py

@@ -4,10 +4,10 @@
 
 import xmltodict
 
-from murfey.client.context import Context
+from murfey.client.context import Context, ensure_dcg_exists
 from murfey.client.contexts.spa import _file_transferred_to, _get_source
-from murfey.client.instance_environment import MurfeyInstanceEnvironment, SampleInfo
-from murfey.util.client import capture_post, get_machine_config_client
+from murfey.client.instance_environment import MurfeyInstanceEnvironment
+from murfey.util.client import capture_post
 from murfey.util.spa_metadata import (
     FoilHoleInfo,
     get_grid_square_atlas_positions,
@@ -69,29 +69,6 @@ def _foil_hole_positions(xml_path: Path, grid_square: int) -> Dict[str, FoilHole
     return foil_holes
 
 
-def _atlas_destination(
-    environment: MurfeyInstanceEnvironment, source: Path, file_path: Path, token: str
-) -> Path:
-    machine_config = get_machine_config_client(
-        str(environment.url.geturl()),
-        token,
-        instrument_name=environment.instrument_name,
-        demo=environment.demo,
-    )
-    for i, destination_part in enumerate(
-        Path(environment.default_destinations[source]).parts
-    ):
-        if destination_part == environment.visit:
-            return Path(machine_config.get("rsync_basepath", "")) / "/".join(
-                Path(environment.default_destinations[source]).parent.parts[: i + 1]
-            )
-    return (
-        Path(machine_config.get("rsync_basepath", ""))
-        / Path(environment.default_destinations[source]).parent
-        / environment.visit
-    )
-
-
 class SPAMetadataContext(Context):
     def __init__(self, acquisition_software: str, basepath: Path, token: str):
         super().__init__("SPA_metadata", acquisition_software, token)
@@ -124,82 +101,19 @@ def post_transfer(
             source = _get_source(transferred_file, environment)
             if not source:
                 logger.warning(
-                    f"Source could not be indentified for {str(transferred_file)}"
+                    f"Source could not be identified for {str(transferred_file)}"
                 )
                 return
 
-            source_visit_dir = source.parent
-
-            logger.info(
-                f"Looking for atlas XML file in metadata directory {str((source_visit_dir / partial_path).parent)}"
-            )
-            atlas_xml_path = list(
-                (source_visit_dir / partial_path).parent.glob("Atlas_*.xml")
-            )[0]
-            logger.info(f"Atlas XML path {str(atlas_xml_path)} found")
-            with open(atlas_xml_path, "rb") as atlas_xml:
-                atlas_xml_data = xmltodict.parse(atlas_xml)
-                atlas_original_pixel_size = float(
-                    atlas_xml_data["MicroscopeImage"]["SpatialScale"]["pixelSize"]["x"][
-                        "numericValue"
-                    ]
-                )
-
-            # need to calculate the pixel size of the downscaled image
-            atlas_pixel_size = atlas_original_pixel_size * 7.8
-            logger.info(f"Atlas image pixel size determined to be {atlas_pixel_size}")
-
-            for p in partial_path.split("/"):
-                if p.startswith("Sample"):
-                    sample = int(p.replace("Sample", ""))
-                    break
-            else:
-                logger.warning(f"Sample could not be identified for {transferred_file}")
-                return
             if source:
-                environment.samples[source] = SampleInfo(
-                    atlas=Path(partial_path), sample=sample
-                )
-                dcg_search_dir = "/".join(
-                    p for p in transferred_file.parent.parts if p != environment.visit
-                )
-                dcg_search_dir = (
-                    dcg_search_dir[1:]
-                    if dcg_search_dir.startswith("//")
-                    else dcg_search_dir
-                )
-                dcg_images_dirs = sorted(
-                    Path(dcg_search_dir).glob("Images-Disc*"),
-                    key=lambda x: x.stat().st_ctime,
-                )
-                if not dcg_images_dirs:
-                    logger.warning(f"Cannot find Images-Disc* in {dcg_search_dir}")
-                    return
-                dcg_tag = str(dcg_images_dirs[-1])
-                dcg_data = {
-                    "experiment_type_id": 37,  # Single particle
-                    "tag": dcg_tag,
-                    "atlas": str(
-                        _atlas_destination(
-                            environment, source, transferred_file, self._token
-                        )
-                        / environment.samples[source].atlas.parent
-                        / atlas_xml_path.with_suffix(".jpg").name
-                    ),
-                    "sample": environment.samples[source].sample,
-                    "atlas_pixel_size": atlas_pixel_size,
-                }
-                capture_post(
-                    base_url=str(environment.url.geturl()),
-                    router_name="workflow.router",
-                    function_name="register_dc_group",
+                dcg_tag = ensure_dcg_exists(
+                    collection_type="spa",
+                    metadata_source=source,
+                    environment=environment,
                     token=self._token,
-                    visit_name=environment.visit,
-                    session_id=environment.murfey_session,
-                    data=dcg_data,
                 )
                 gs_pix_positions = get_grid_square_atlas_positions(
-                    source_visit_dir / partial_path
+                    source.parent / partial_path
                 )
                 for gs, pos_data in gs_pix_positions.items():
                     if pos_data:
@@ -228,46 +142,21 @@ def post_transfer(
             and environment
         ):
             # Make sure we have a data collection group before trying to register grid square
-            dcg_search_dir = "/".join(
-                p
-                for p in transferred_file.parent.parent.parts
-                if p != environment.visit
-            )
-            dcg_search_dir = (
-                dcg_search_dir[1:]
-                if dcg_search_dir.startswith("//")
-                else dcg_search_dir
-            )
-            dcg_images_dirs = sorted(
-                Path(dcg_search_dir).glob("Images-Disc*"),
-                key=lambda x: x.stat().st_ctime,
-            )
-            if not dcg_images_dirs:
-                logger.warning(f"Cannot find Images-Disc* in {dcg_search_dir}")
-                return
-            dcg_tag = str(dcg_images_dirs[-1])
-            dcg_data = {
-                "experiment_type_id": 37,  # Single particle
-                "tag": dcg_tag,
-            }
-            capture_post(
-                base_url=str(environment.url.geturl()),
-                router_name="workflow.router",
-                function_name="register_dc_group",
+            source = _get_source(transferred_file, environment=environment)
+            if source is None:
+                return None
+            ensure_dcg_exists(
+                collection_type="spa",
+                metadata_source=source,
+                environment=environment,
                 token=self._token,
-                visit_name=environment.visit,
-                session_id=environment.murfey_session,
-                data=dcg_data,
             )
 
             gs_name = int(transferred_file.stem.split("_")[1])
             logger.info(
                 f"Collecting foil hole positions for {str(transferred_file)} and grid square {gs_name}"
             )
             fh_positions = _foil_hole_positions(transferred_file, gs_name)
-            source = _get_source(transferred_file, environment=environment)
-            if source is None:
-                return None
             visitless_source_search_dir = str(source).replace(
                 f"/{environment.visit}", ""
             )