Register positions of tomograms on search maps

stephen-riggs · stephen-riggs · commit 405d259ec0a2 · 2025-06-26T16:50:11.000+01:00
diff --git a/src/murfey/client/contexts/tomo_metadata.py b/src/murfey/client/contexts/tomo_metadata.py
@@ -277,7 +277,9 @@ def post_transfer(
                         "tag": visitless_source,
                         "x_stage_position": batch_stage_location_x,
                         "y_stage_position": batch_stage_location_y,
-                        "search_map": search_map_name,
+                        "x_beamshift": 0,
+                        "y_beamshift": 0,
+                        "search_map_name": search_map_name,
                     },
                 )
 
@@ -298,8 +300,10 @@ def post_transfer(
                             bp_url,
                             json={
                                 "tag": visitless_source,
-                                "x_stage_position": beamshift_position_x,
-                                "y_stage_position": beamshift_position_y,
-                                "search_map": search_map_name,
+                                "x_stage_position": batch_stage_location_x,
+                                "y_stage_position": batch_stage_location_y,
+                                "x_beamshift": beamshift_position_x,
+                                "y_beamshift": beamshift_position_y,
+                                "search_map_name": search_map_name,
                             },
                         )
diff --git a/src/murfey/server/api/workflow.py b/src/murfey/server/api/workflow.py
@@ -819,6 +819,9 @@ def register_completed_tilt_series(
                     "pixel_size": preproc_params.pixel_size,
                     "manual_tilt_offset": -tilt_offset,
                     "node_creator_queue": machine_config.node_creator_queue,
+                    "search_map_id": ts.search_map_id,
+                    "x_location": ts.x_location,
+                    "y_location": ts.y_location,
                 },
             }
             if _transport_object:
diff --git a/src/murfey/server/feedback.py b/src/murfey/server/feedback.py
@@ -1937,6 +1937,9 @@ def feedback_callback(header: dict, message: dict) -> None:
                         "pixel_size": preproc_params.pixel_size,
                         "manual_tilt_offset": -tilt_offset,
                         "node_creator_queue": machine_config.node_creator_queue,
+                        "search_map_id": relevant_tilt_series.search_map_id,
+                        "x_location": relevant_tilt_series.x_location,
+                        "y_location": relevant_tilt_series.y_location,
                     },
                 }
                 if murfey.server._transport_object:
diff --git a/src/murfey/util/db.py b/src/murfey/util/db.py
@@ -349,15 +349,20 @@ class SessionProcessingParameters(SQLModel, table=True):  # type: ignore
 
 class TiltSeries(SQLModel, table=True):  # type: ignore
     id: int = Field(primary_key=True)
+    ispyb_id: Optional[int] = None
     tag: str
     rsync_source: str
     session_id: int = Field(foreign_key="session.id")
+    search_map_id: int = Field(foreign_key="searchmap.id")
     tilt_series_length: int = -1
     processing_requested: bool = False
+    x_location: Optional[float] = None
+    y_location: Optional[float] = None
     session: Optional[Session] = Relationship(back_populates="tilt_series")
     tilts: List["Tilt"] = Relationship(
         back_populates="tilt_series", sa_relationship_kwargs={"cascade": "delete"}
     )
+    search_map: Optional["GridSquare"] = Relationship(back_populates="batch_positions")
 
 
 class Tilt(SQLModel, table=True):  # type: ignore
@@ -620,7 +625,12 @@ class SearchMap(SQLModel, table=True):  # type: ignore
     reference_matrix: Optional[dict[str, float]] = None
     stage_correction: Optional[dict[str, float]] = None
     image_shift_correction: Optional[dict[str, float]] = None
+    width: Optional[int] = None
+    height: Optional[int] = None
     session: Optional[Session] = Relationship(back_populates="search_maps")
+    batch_positions: List["TiltSeries"] = Relationship(
+        back_populates="search_map", sa_relationship_kwargs={"cascade": "delete"}
+    )
 
 
 class Movie(SQLModel, table=True):  # type: ignore
diff --git a/src/murfey/util/models.py b/src/murfey/util/models.py
@@ -157,8 +157,10 @@ class SearchMapParameters:
 class BatchPositionParameters:
     tag: str
     x_stage_position: float
-    Y_stage_position: float
-    search_map: str
+    y_stage_position: float
+    x_beamshift: float
+    y_beamshift: float
+    search_map_name: str
 
 
 class MultigridWatcherSetup(BaseModel):
diff --git a/src/murfey/util/tomo_metadata.py b/src/murfey/util/tomo_metadata.py
@@ -1,3 +1,5 @@
+import logging
+
 import numpy as np
 from sqlmodel import Session, select
 
@@ -7,8 +9,11 @@
 from murfey.util.config import get_machine_config
 from murfey.util.db import DataCollectionGroup, SearchMap
 from murfey.util.db import Session as MurfeySession
+from murfey.util.db import TiltSeries
 from murfey.util.models import BatchPositionParameters, SearchMapParameters
 
+logger = logging.getLogger("murfey.client.util.tomo_metadata")
+
 
 def register_search_map_in_database(
     session_id: MurfeySessionID,
@@ -127,33 +132,22 @@ def register_search_map_in_database(
                 ],
             ]
         )
-        M_corrected = np.matmul(np.linalg.inv(M), R)
-        vector_pixel = np.matmul(M_corrected, B)
+        vector_pixel = np.matmul(np.linalg.inv(M), np.matmul(R, np.matmul(M, B)))
 
         camera = getattr(Camera, machine_config.camera)
-        if camera == Camera.FALCON:
-            vector_pixel = np.matmul(np.array([[-1, 0], [0, -1]]), vector_pixel)
+        if camera == Camera.FALCON or Camera.K3_FLIPY:
+            vector_pixel = np.matmul(np.array([[1, 0], [0, -1]]), vector_pixel)
+        elif camera == Camera.K3_FLIPX:
+            vector_pixel = np.matmul(np.array([[-1, 0], [0, 1]]), vector_pixel)
 
         search_map_params.height_on_atlas = int(
             search_map.height * search_map.pixel_size / dcg.atlas_pixel_size
         )
         search_map_params.width_on_atlas = int(
             search_map.width * search_map.pixel_size / dcg.atlas_pixel_size
         )
-        search_map_params.x_location = (
-            vector_pixel[0]
-            / dcg.atlas_binning
-            * search_map.pixel_size
-            / dcg.atlas_pixel_size
-            + 2003
-        )
-        search_map_params.y_location = (
-            vector_pixel[1]
-            / dcg.atlas_binning
-            * search_map.pixel_size
-            / dcg.atlas_pixel_size
-            + 2003
-        )
+        search_map_params.x_location = vector_pixel[0] / dcg.atlas_pixel_size + 2003
+        search_map_params.y_location = vector_pixel[1] / dcg.atlas_pixel_size + 2003
         search_map.x_location = search_map_params.x_location
         search_map.y_location = search_map_params.y_location
         if _transport_object:
@@ -169,4 +163,86 @@ def register_batch_position_in_database(
     batch_parameters: BatchPositionParameters,
     murfey_db: Session,
 ):
-    pass
+    search_map = murfey_db.exec(
+        select(SearchMap)
+        .where(SearchMap.name == batch_parameters.search_map_name)
+        .where(SearchMap.tag == batch_parameters.tag)
+        .where(SearchMap.session_id == session_id)
+    ).one()
+
+    try:
+        tilt_series = murfey_db.exec(
+            select(TiltSeries)
+            .where(TiltSeries.tag == batch_name)
+            .where(TiltSeries.session_id == session_id)
+        ).one()
+        if tilt_series.x_location:
+            logger.info(f"Already did position analysis for tomogram {batch_name}")
+            return
+    except Exception:
+        tilt_series = TiltSeries(
+            tag=batch_name,
+            rsync_source=batch_parameters.tag,
+            session_id=session_id,
+            search_map_id=search_map.id,
+        )
+
+    # Get the pixel location on the searchmap
+    M = np.array(
+        [
+            [
+                search_map.reference_matrix["m11"],
+                search_map.reference_matrix["m12"],
+            ],
+            [
+                search_map.reference_matrix["m21"],
+                search_map.reference_matrix["m22"],
+            ],
+        ]
+    )
+    R1 = np.array(
+        [
+            [
+                search_map.stage_correction["m11"],
+                search_map.stage_correction["m12"],
+            ],
+            [
+                search_map.stage_correction["m21"],
+                search_map.stage_correction["m22"],
+            ],
+        ]
+    )
+    R2 = np.array(
+        [
+            [
+                search_map.image_shift_correction["m11"],
+                search_map.image_shift_correction["m12"],
+            ],
+            [
+                search_map.image_shift_correction["m21"],
+                search_map.image_shift_correction["m22"],
+            ],
+        ]
+    )
+
+    A = np.array([search_map.x_stage_position, search_map.y_stage_position])
+    B = np.array([batch_parameters.x_stage_position, batch_parameters.y_stage_position])
+
+    vector_pixel = np.matmul(
+        np.linalg.inv(M),
+        np.matmul(np.linalg.inv(R1), np.matmul(np.linalg.inv(R2), np.matmul(M, B - A))),
+    )
+    centre_batch_pixel = vector_pixel / search_map.pixel_size + [
+        search_map.width / 2,
+        search_map.height / 2,
+    ]
+    tilt_series.x_location = (
+        centre_batch_pixel[0]
+        - BatchPositionParameters.x_beamshift / search_map.pixel_size
+    )
+    tilt_series.y_location = (
+        centre_batch_pixel[1]
+        - BatchPositionParameters.y_beamshift / search_map.pixel_size
+    )
+    murfey_db.add(tilt_series)
+    murfey_db.commit()
