Need frame counts for tomography dose weighting (#469)

Adds frame count finding for tomography, and passes this on to the processing.
Storing the frame count requires a database change, so we may not want to merge this change immediately.

Author: stephen-riggs

src/murfey/client/contexts/tomo.py

@@ -78,6 +78,7 @@ class TomographyContext(Context):
         ProcessingParameter("image_size_y", "Image Size Y"),
         ProcessingParameter("pixel_size_on_image", "Pixel Size"),
         ProcessingParameter("motion_corr_binning", "Motion Correction Binning"),
+        ProcessingParameter("frame_count", "Number of image frames"),
         ProcessingParameter("num_eer_frames", "Number of EER Frames"),
     ]
 
@@ -361,6 +362,9 @@ def _add_tilt(
                 "dose_per_frame": environment.data_collection_parameters.get(
                     "dose_per_frame", 0
                 ),
+                "frame_count": environment.data_collection_parameters.get(
+                    "frame_count", 0
+                ),
                 "mc_binning": environment.data_collection_parameters.get(
                     "motion_corr_binning", 1
                 ),
@@ -591,6 +595,7 @@ def gather_metadata(
             mdoc_metadata: OrderedDict = OrderedDict({})
             mdoc_metadata["experiment_type"] = "tomography"
             mdoc_metadata["voltage"] = float(mdoc_data["Voltage"])
+            mdoc_metadata["frame_count"] = int(mdoc_data_block["NumSubFrames"])
             mdoc_metadata["image_size_x"] = int(mdoc_data["ImageSize"][0])
             mdoc_metadata["image_size_y"] = int(mdoc_data["ImageSize"][1])
             mdoc_metadata["magnification"] = int(mdoc_data_block["Magnification"])
@@ -658,6 +663,9 @@ def gather_metadata(
                 mdoc_metadata["num_eer_frames"] = murfey.util.eer.num_frames(
                     metadata_file.parent / data_file
                 )
+                mdoc_metadata["frame_count"] = int(
+                    mdoc_metadata["eer_fractionation"] / mdoc_metadata["num_eer_frames"]
+                )
         except Exception as e:
             logger.error(f"Exception encountered in metadata gathering: {str(e)}")
             return OrderedDict({})

src/murfey/server/__init__.py

@@ -1951,11 +1951,7 @@ def _flush_tomography_preprocessing(message: dict):
         .where(db.DataCollectionGroup.session_id == session_id)
         .where(db.DataCollectionGroup.tag == message["data_collection_group_tag"])
     ).first()
-    proc_params = murfey_db.exec(
-        select(db.TomographyPreprocessingParameters).where(
-            db.TomographyPreprocessingParameters.dcg_id == collected_ids.id
-        )
-    ).one()
+    proc_params = get_tomo_preproc_params(collected_ids.id)
     if not proc_params:
         visit_name = message["visit_name"].replace("\r\n", "").replace("\n", "")
         logger.warning(
@@ -2009,6 +2005,7 @@ def _flush_tomography_preprocessing(message: dict):
                 "mc_uuid": murfey_ids[0],
                 "ft_bin": proc_params.motion_corr_binning,
                 "fm_dose": proc_params.dose_per_frame,
+                "frame_count": proc_params.frame_count,
                 "gain_ref": (
                     str(machine_config.rsync_basepath / proc_params.gain_ref)
                     if proc_params.gain_ref
@@ -2471,6 +2468,9 @@ def feedback_callback(header: dict, message: dict) -> None:
                         "dcid": ids.dcid,
                         "appid": ids.appid,
                         "stack_file": str(stack_file),
+                        "dose_per_frame": preproc_params.dose_per_frame,
+                        "frame_count": preproc_params.frame_count,
+                        "kv": preproc_params.voltage,
                         "pixel_size": preproc_params.pixel_size,
                         "manual_tilt_offset": -tilt_offset,
                         "node_creator_queue": machine_config.node_creator_queue,
@@ -2829,6 +2829,7 @@ def feedback_callback(header: dict, message: dict) -> None:
                     pixel_size=float(message["pixel_size_on_image"]) * 10**10,
                     voltage=message["voltage"],
                     dose_per_frame=message["dose_per_frame"],
+                    frame_count=message["frame_count"],
                     motion_corr_binning=message["motion_corr_binning"],
                     gain_ref=message["gain_ref"],
                     eer_fractionation_file=message["eer_fractionation_file"],

src/murfey/server/api/__init__.py

@@ -729,6 +729,9 @@ def register_completed_tilt_series(
                     "dcid": ids.dcid,
                     "appid": ids.appid,
                     "stack_file": str(stack_file),
+                    "dose_per_frame": preproc_params.dose_per_frame,
+                    "frame_count": preproc_params.frame_count,
+                    "kv": preproc_params.voltage,
                     "pixel_size": preproc_params.pixel_size,
                     "manual_tilt_offset": -tilt_offset,
                     "node_creator_queue": machine_config.node_creator_queue,
@@ -1120,6 +1123,9 @@ async def request_tomography_preprocessing(
         / str(ppath.stem + "_motion_corrected.mrc")
     )
     mrc_out = Path("/".join(secure_filename(p) for p in mrc_out.parts))
+
+    recipe_name = machine_config.recipes.get("em-tomo-preprocess", "em-tomo-preprocess")
+
     data_collection = db.exec(
         select(DataCollectionGroup, DataCollection, ProcessingJob, AutoProcProgram)
         .where(DataCollectionGroup.session_id == session_id)
@@ -1128,7 +1134,7 @@ async def request_tomography_preprocessing(
         .where(DataCollection.dcg_id == DataCollectionGroup.id)
         .where(ProcessingJob.dc_id == DataCollection.id)
         .where(AutoProcProgram.pj_id == ProcessingJob.id)
-        .where(ProcessingJob.recipe == "em-tomo-preprocess")
+        .where(ProcessingJob.recipe == recipe_name)
     ).all()
     if data_collection:
         if registered_tilts := db.exec(
@@ -1143,7 +1149,7 @@ async def request_tomography_preprocessing(
         if not mrc_out.parent.exists():
             mrc_out.parent.mkdir(parents=True, exist_ok=True)
         zocalo_message: dict = {
-            "recipes": ["em-tomo-preprocess"],
+            "recipes": [recipe_name],
             "parameters": {
                 "node_creator_queue": machine_config.node_creator_queue,
                 "dcid": dcid,
@@ -1158,6 +1164,7 @@ async def request_tomography_preprocessing(
                 "mc_uuid": murfey_ids[0],
                 "ft_bin": proc_file.mc_binning,
                 "fm_dose": proc_file.dose_per_frame,
+                "frame_count": proc_file.frame_count,
                 "gain_ref": (
                     str(machine_config.rsync_basepath / proc_file.gain_ref)
                     if proc_file.gain_ref and machine_config.data_transfer_enabled

src/murfey/util/db.py

@@ -458,6 +458,7 @@ class TomographyPreprocessingParameters(SQLModel, table=True):  # type: ignore
     dcg_id: int = Field(primary_key=True, foreign_key="datacollectiongroup.id")
     pixel_size: float
     dose_per_frame: float
+    frame_count: int
     voltage: int
     eer_fractionation_file: Optional[str] = None
     motion_corr_binning: int = 1

src/murfey/util/models.py

@@ -349,6 +349,7 @@ class CompletedTiltSeries(BaseModel):
 
 class PreprocessingParametersTomo(BaseModel):
     dose_per_frame: float
+    frame_count: int
     gain_ref: Optional[str]
     experiment_type: str
     voltage: float