Search for user provided initial model (#312)

Appropriately rescale if found

Author: d-j-hatton

src/murfey/server/__init__.py

@@ -2,7 +2,9 @@
 
 import argparse
 import logging
+import math
 import os
+import subprocess
 import time
 from datetime import datetime
 from functools import partial, singledispatch, wraps
@@ -39,7 +41,12 @@
 import murfey.server.prometheus as prom
 import murfey.server.websocket
 from murfey.client.contexts.tomo import _midpoint
-from murfey.server.config import get_hostname, get_machine_config, get_microscope
+from murfey.server.config import (
+    MachineConfig,
+    get_hostname,
+    get_machine_config,
+    get_microscope,
+)
 from murfey.server.murfey_db import url  # murfey_db
 
 try:
@@ -1547,6 +1554,95 @@ def _register_class_selection(message: dict, _db=murfey_db, demo: bool = False):
     _db.close()
 
 
+def _find_initial_model(visit: str, machine_config: MachineConfig) -> Path | None:
+    if machine_config.initial_model_search_directory:
+        visit_directory = (
+            machine_config.rsync_basepath
+            / (machine_config.rsync_module or "data")
+            / str(datetime.now().year)
+            / visit
+        )
+        possible_models = list(
+            (visit_directory / machine_config.initial_model_search_directory).glob(
+                "*.mrc"
+            )
+        )
+        if possible_models:
+            return sorted(possible_models, key=lambda x: x.stat().st_ctime)[-1]
+    return None
+
+
+def _downscaled_box_size(
+    particle_diameter: int, pixel_size: float
+) -> Tuple[int, float]:
+    box_size = int(math.ceil(1.2 * particle_diameter))
+    box_size = box_size + box_size % 2
+    for small_box_pix in (
+        48,
+        64,
+        96,
+        128,
+        160,
+        192,
+        256,
+        288,
+        300,
+        320,
+        360,
+        384,
+        400,
+        420,
+        450,
+        480,
+        512,
+        640,
+        768,
+        896,
+        1024,
+    ):
+        # Don't go larger than the original box
+        if small_box_pix > box_size:
+            return box_size, pixel_size
+        # If Nyquist freq. is better than 8.5 A, use this downscaled box, else step size
+        small_box_angpix = pixel_size * box_size / small_box_pix
+        if small_box_angpix < 4.25:
+            return small_box_pix, small_box_angpix
+    raise ValueError(f"Box size is too large: {box_size}")
+
+
+def _resize_intial_model(
+    downscaled_box_size: int,
+    downscaled_pixel_size: float,
+    input_path: Path,
+    output_path: Path,
+    executables: Dict[str, str],
+    env: Dict[str, str],
+) -> None:
+    if executables.get("relion_image_handler"):
+        comp_proc = subprocess.run(
+            [
+                f"{executables['relion_image_handler']}",
+                "--i",
+                str(input_path),
+                "--new_box",
+                str(downscaled_box_size),
+                "--rescale_angpix",
+                str(downscaled_pixel_size),
+                "--force_header_angpix",
+                "--o",
+                str(output_path),
+            ],
+            capture_output=True,
+            text=True,
+            env=env,
+        )
+        if comp_proc.returncode:
+            logger.error(
+                f"Resizing initial model {input_path} failed \n {comp_proc.stdout}"
+            )
+    return None
+
+
 def _register_3d_batch(message: dict, _db=murfey_db, demo: bool = False):
     """Received 3d batch from class selection service"""
     class3d_message = message.get("class3d_message")
@@ -1567,6 +1663,36 @@ def _register_3d_batch(message: dict, _db=murfey_db, demo: bool = False):
     ).one()
     other_options = dict(feedback_params)
 
+    visit_name = (
+        _db.exec(
+            select(db.ClientEnvironment).where(
+                db.ClientEnvironment.session_id == message["session_id"]
+            )
+        )
+        .one()
+        .visit
+    )
+
+    provided_initial_model = _find_initial_model(visit_name, machine_config)
+    if provided_initial_model:
+        rescaled_initial_model_path = (
+            provided_initial_model.parent
+            / f"{provided_initial_model.stem}_rescaled{provided_initial_model.suffix}"
+        )
+        _resize_intial_model(
+            *_downscaled_box_size(
+                message["particle_diameter"],
+                relion_options["angpix"],
+            ),
+            provided_initial_model,
+            rescaled_initial_model_path,
+            machine_config.external_executables,
+            machine_config.external_environment,
+        )
+        feedback_params.initial_model = str(rescaled_initial_model_path)
+        _db.add(feedback_params)
+        _db.commit()
+
     if feedback_params.hold_class3d:
         # If waiting then save the message
         class3d_params = _db.exec(

src/murfey/server/config.py

@@ -58,6 +58,7 @@ class MachineConfig(BaseModel):
     upstream_data_tiff_locations: List[str] = ["processed"]  # Location of CLEM TIFFs
 
     model_search_directory: str = "processing"
+    initial_model_search_directory: str = "processing/initial_model"
 
     failure_queue: str = ""
     auth_key: str = ""