Add ISPyB inserts for grid squares

Author: d-j-hatton

src/murfey/client/contexts/spa.py

@@ -59,9 +59,18 @@ class GridSquare(NamedTuple):
     tag: str = ""
 
 
-def _get_grid_square_atlas_positions(
-    xml_path: Path, grid_square: str = ""
-) -> Dict[str, Tuple[Optional[int], Optional[int], Optional[float], Optional[float]]]:
+def _get_grid_square_atlas_positions(xml_path: Path, grid_square: str = "") -> Dict[
+    str,
+    Tuple[
+        Optional[int],
+        Optional[int],
+        Optional[float],
+        Optional[float],
+        Optional[int],
+        Optional[int],
+        Optional[float],
+    ],
+]:
     with open(
         xml_path,
         "r",
@@ -71,7 +80,16 @@ def _get_grid_square_atlas_positions(
         "TileXml"
     ]
     gs_pix_positions: Dict[
-        str, Tuple[Optional[int], Optional[int], Optional[float], Optional[float]]
+        str,
+        Tuple[
+            Optional[int],
+            Optional[int],
+            Optional[float],
+            Optional[float],
+            Optional[int],
+            Optional[int],
+            Optional[float],
+        ],
     ] = {}
     for ti in tile_info:
         try:
@@ -96,6 +114,13 @@ def _get_grid_square_atlas_positions(
                     * 1e9,
                     float(gs["value"]["b:PositionOnTheAtlas"]["c:Physical"]["d:y"])
                     * 1e9,
+                    int(
+                        float(gs["value"]["b:PositionOnTheAtlas"]["c:Size"]["d:width"])
+                    ),
+                    int(
+                        float(gs["value"]["b:PositionOnTheAtlas"]["c:Size"]["d:height"])
+                    ),
+                    float(gs["value"]["b:PositionOnTheAtlas"]["c:Rotation"]),
                 )
                 if grid_square:
                     break
@@ -557,7 +582,10 @@ def _position_analysis(
                 Optional[int],
                 Optional[float],
                 Optional[float],
-            ] = (None, None, None, None)
+                Optional[int],
+                Optional[int],
+                Optional[float],
+            ] = (None, None, None, None, None, None, None)
             data_collection_group = (
                 requests.get(
                     f"{str(environment.url.geturl())}/sessions/{environment.murfey_session}/data_collection_groups"
@@ -611,6 +639,9 @@ def _position_analysis(
                     "y_location": gs_pix_position[1],
                     "x_stage_position": gs_pix_position[2],
                     "y_stage_position": gs_pix_position[3],
+                    "width": gs_pix_position[4],
+                    "height": gs_pix_position[5],
+                    "angle": gs_pix_position[6],
                 },
             )
         foil_hole = _foil_hole_from_file(transferred_file)

src/murfey/client/contexts/spa_metadata.py

@@ -155,5 +155,8 @@ def post_transfer(
                                     "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],
                                 },
                             )

src/murfey/server/api/__init__.py

@@ -491,6 +491,13 @@ def register_grid_square(
     grid_square_params: GridSquareParameters,
     db=murfey_db,
 ):
+    if _transport_object:
+        dcg = db.exec(
+            select(DataCollectionGroup)
+            .where(DataCollectionGroup.session_id == session_id)
+            .where(DataCollectionGroup.tag == grid_square_params.tag)
+        ).one()
+        _transport_object.do_insert_grid_square(dcg.atlas_id, gsid, grid_square_params)
     try:
         grid_square = db.exec(
             select(GridSquare)

src/murfey/server/ispyb.py

@@ -22,14 +22,15 @@
     BLSubSample,
     DataCollection,
     DataCollectionGroup,
+    GridSquare,
     ProcessingJob,
     ProcessingJobParameter,
     Proposal,
     ZcZocaloBuffer,
     url,
 )
 
-from murfey.util.models import Sample, Visit
+from murfey.util.models import GridSquareParameters, Sample, Visit
 
 log = logging.getLogger("murfey.server.ispyb")
 
@@ -127,6 +128,39 @@ def do_update_atlas(
             )
         return {"success": False, "return_value": None}
 
+    def do_insert_grid_square(
+        self,
+        atlas_id: int,
+        grid_square_id: int,
+        grid_square_parameters: GridSquareParameters,
+    ):
+        record = GridSquare(
+            atlasId=atlas_id,
+            gridSquareLabel=grid_square_id,
+            gridSquareImage=grid_square_parameters.image,
+            pixelLocationX=grid_square_parameters.x_location,
+            pixelLocationY=grid_square_parameters.y_location,
+            height=grid_square_parameters.height,
+            weight=grid_square_parameters.width,
+            angle=grid_square_parameters.angle,
+            stageLocationX=grid_square_parameters.x_stage_position,
+            stageLocationY=grid_square_parameters.y_stage_position,
+            pixelSize=grid_square_parameters.pixel_size,
+        )
+        try:
+            with Session() as db:
+                db.add(record)
+                db.commit()
+                log.info(f"Created GridSquare {record.gridSquareId}")
+                return {"success": True, "return_value": record.gridSquareId}
+        except ispyb.ISPyBException as e:
+            log.error(
+                "Inserting GridSquare entry caused exception '%s'.",
+                e,
+                exc_info=True,
+            )
+        return {"success": False, "return_value": None}
+
     def send(self, queue: str, message: dict, new_connection: bool = False):
         if self.transport:
             if not self.transport.is_connected():

src/murfey/util/models.py

@@ -273,8 +273,11 @@ class GridSquareParameters(BaseModel):
     readout_area_y: Optional[int] = None
     thumbnail_size_x: Optional[int] = None
     thumbnail_size_y: Optional[int] = None
+    height: Optional[int] = None
+    width: Optional[int] = None
     pixel_size: Optional[float] = None
     image: str = ""
+    angle: Optional[float] = None
 
 
 class FoilHoleParameters(BaseModel):