Resolved merge conflicts with 'main' branch

Author: tieneupin

src/murfey/server/api/spa.py

@@ -27,7 +27,6 @@
 import murfey.server.prometheus as prom
 from murfey.server import _transport_object
 from murfey.server.api.auth import MurfeySessionID, validate_instrument_token
-from murfey.server.api.spa import _cryolo_model_path
 from murfey.server.feedback import (
     _murfey_id,
     check_tilt_series_mc,
@@ -58,7 +57,11 @@
     TiltSeries,
 )
 from murfey.util.models import ProcessingParametersSPA, ProcessingParametersTomo
-from murfey.util.processing_params import default_spa_parameters
+from murfey.util.processing_params import (
+    cryolo_model_path,
+    default_spa_parameters,
+    motion_corrected_mrc,
+)
 from murfey.util.tomo import midpoint
 
 logger = getLogger("murfey.server.api.workflow")
@@ -365,29 +368,7 @@ async def request_spa_preprocessing(
     machine_config = get_machine_config(instrument_name=instrument_name)[
         instrument_name
     ]
-    parts = [secure_filename(p) for p in Path(proc_file.path).parts]
-    visit_idx = parts.index(visit_name)
-    core = Path("/") / Path(*parts[: visit_idx + 1])
-    ppath = Path("/") / Path(*parts)
-    sub_dataset = ppath.relative_to(core).parts[0]
-    extra_path = machine_config.processed_extra_directory
-    for i, p in enumerate(ppath.parts):
-        if p.startswith("raw"):
-            movies_path_index = i
-            break
-    else:
-        raise ValueError(f"{proc_file.path} does not contain a raw directory")
-    mrc_out = (
-        core
-        / machine_config.processed_directory_name
-        / sub_dataset
-        / extra_path
-        / "MotionCorr"
-        / "job002"
-        / "Movies"
-        / "/".join(ppath.parts[movies_path_index + 1 : -1])
-        / str(ppath.stem + "_motion_corrected.mrc")
-    )
+    mrc_out = motion_corrected_mrc(Path(proc_file.path), visit_name, machine_config)
     try:
         collected_ids = db.exec(
             select(DataCollectionGroup, DataCollection, ProcessingJob, AutoProcProgram)
@@ -479,7 +460,7 @@ async def request_spa_preprocessing(
                 ),
                 "do_icebreaker_jobs": default_spa_parameters.do_icebreaker_jobs,
                 "cryolo_model_weights": str(
-                    _cryolo_model_path(visit_name, instrument_name)
+                    cryolo_model_path(visit_name, instrument_name)
                 ),
             },
         }
@@ -491,7 +472,8 @@ async def request_spa_preprocessing(
             _transport_object.send("processing_recipe", zocalo_message)
         else:
             logger.error(
-                f"Pe-processing was requested for {sanitise(ppath.name)} but no Zocalo transport object was found"
+                f"Pre-processing was requested for {sanitise(Path(proc_file.path).name)} "
+                "but no Zocalo transport object was found"
             )
             return proc_file
 
@@ -646,22 +628,7 @@ async def request_tomography_preprocessing(
     machine_config = get_machine_config(instrument_name=instrument_name)[
         instrument_name
     ]
-    visit_idx = Path(proc_file.path).parts.index(visit_name)
-    core = Path(*Path(proc_file.path).parts[: visit_idx + 1])
-    ppath = Path("/".join(secure_filename(p) for p in Path(proc_file.path).parts))
-    sub_dataset = "/".join(ppath.relative_to(core).parts[:-1])
-    extra_path = machine_config.processed_extra_directory
-    mrc_out = (
-        core
-        / machine_config.processed_directory_name
-        / sub_dataset
-        / extra_path
-        / "MotionCorr"
-        / "job002"
-        / "Movies"
-        / str(ppath.stem + "_motion_corrected.mrc")
-    )
-    mrc_out = Path("/".join(secure_filename(p) for p in mrc_out.parts))
+    mrc_out = motion_corrected_mrc(Path(proc_file.path), visit_name, machine_config)
 
     recipe_name = machine_config.recipes.get("em-tomo-preprocess", "em-tomo-preprocess")
 
@@ -732,7 +699,8 @@ async def request_tomography_preprocessing(
             _transport_object.send("processing_recipe", zocalo_message)
         else:
             logger.error(
-                f"Pe-processing was requested for {sanitise(ppath.name)} but no Zocalo transport object was found"
+                f"Pre-processing was requested for {sanitise(Path(proc_file.path).name)} "
+                f"but no Zocalo transport object was found"
             )
             return proc_file
     else:

src/murfey/server/api/workflow.py

@@ -51,7 +51,7 @@
     get_microscope,
     get_security_config,
 )
-from murfey.util.processing_params import default_spa_parameters
+from murfey.util.processing_params import default_spa_parameters, motion_corrected_mrc
 from murfey.util.tomo import midpoint
 
 logger = logging.getLogger("murfey.server.feedback")
@@ -106,37 +106,16 @@ def get_all_tilts(tilt_series_id: int) -> List[str]:
     ).all()
     if not complete_results:
         return []
+    visit_name = complete_results[0][2].visit
     instrument_name = complete_results[0][2].instrument_name
     results = [r[0] for r in complete_results]
     machine_config = get_machine_config(instrument_name=instrument_name)[
         instrument_name
     ]
-
-    def _mc_path(mov_path: Path) -> str:
-        for p in mov_path.parts:
-            if "-" in p and p.startswith(("bi", "nr", "nt", "cm", "sw")):
-                visit_name = p
-                break
-        else:
-            raise ValueError(f"No visit found in {mov_path}")
-        visit_idx = Path(mov_path).parts.index(visit_name)
-        core = Path(*Path(mov_path).parts[: visit_idx + 1])
-        ppath = Path(mov_path)
-        sub_dataset = "/".join(ppath.relative_to(core).parts[:-1])
-        extra_path = machine_config.processed_extra_directory
-        mrc_out = (
-            core
-            / machine_config.processed_directory_name
-            / sub_dataset
-            / extra_path
-            / "MotionCorr"
-            / "job002"
-            / "Movies"
-            / str(ppath.stem + "_motion_corrected.mrc")
-        )
-        return str(mrc_out)
-
-    return [_mc_path(Path(r.movie_path)) for r in results]
+    return [
+        str(motion_corrected_mrc(Path(r.movie_path), visit_name, machine_config))
+        for r in results
+    ]
 
 
 def get_job_ids(tilt_series_id: int, appid: int) -> JobIDs:

src/murfey/server/feedback.py

@@ -22,7 +22,6 @@
 import murfey.server.api.prometheus
 import murfey.server.api.session_control
 import murfey.server.api.session_info
-import murfey.server.api.spa
 import murfey.server.api.websocket
 import murfey.server.api.workflow
 from murfey.server import template_files
@@ -90,7 +89,6 @@ class Settings(BaseSettings):
 app.include_router(murfey.server.api.workflow.correlative_router)
 app.include_router(murfey.server.api.workflow.spa_router)
 app.include_router(murfey.server.api.workflow.tomo_router)
-app.include_router(murfey.server.api.spa.router)
 app.include_router(murfey.server.api.clem.router)
 
 app.include_router(murfey.server.api.prometheus.router)

src/murfey/server/main.py

@@ -57,6 +57,7 @@ class MachineConfig(BaseModel):  # type: ignore
     processing_enabled: bool = True
     process_by_default: bool = True
     gain_directory_name: str = "processing"
+    process_multiple_datasets: bool = True
     processed_directory_name: str = "processed"
     processed_extra_directory: str = ""
     recipes: dict[str, str] = {

src/murfey/util/config.py

@@ -1,6 +1,55 @@
+from datetime import datetime
+from functools import lru_cache
+from pathlib import Path
 from typing import Literal, Optional
 
 from pydantic import BaseModel
+from werkzeug.utils import secure_filename
+
+from murfey.util.config import MachineConfig, get_machine_config
+
+
+def motion_corrected_mrc(
+    input_movie: Path, visit_name: str, machine_config: MachineConfig
+):
+    parts = [secure_filename(p) for p in input_movie.parts]
+    visit_idx = parts.index(visit_name)
+    core = Path("/") / Path(*parts[: visit_idx + 1])
+    ppath = Path("/") / Path(*parts)
+    if machine_config.process_multiple_datasets:
+        sub_dataset = ppath.relative_to(core).parts[0]
+    else:
+        sub_dataset = ""
+    extra_path = machine_config.processed_extra_directory
+    mrc_out = (
+        core
+        / machine_config.processed_directory_name
+        / sub_dataset
+        / extra_path
+        / "MotionCorr"
+        / "job002"
+        / "Movies"
+        / ppath.parent.relative_to(core / sub_dataset)
+        / str(ppath.stem + "_motion_corrected.mrc")
+    )
+    return Path("/".join(secure_filename(p) for p in mrc_out.parts))
+
+
+@lru_cache(maxsize=5)
+def cryolo_model_path(visit: str, instrument_name: str) -> Path:
+    machine_config = get_machine_config(instrument_name=instrument_name)[
+        instrument_name
+    ]
+    if machine_config.model_search_directory:
+        visit_directory = (
+            machine_config.rsync_basepath / str(datetime.now().year) / visit
+        )
+        possible_models = list(
+            (visit_directory / machine_config.model_search_directory).glob("*.h5")
+        )
+        if possible_models:
+            return sorted(possible_models, key=lambda x: x.stat().st_ctime)[-1]
+    return machine_config.default_model
 
 
 class CLEMAlignAndMergeParameters(BaseModel):

src/murfey/util/processing_params.py

@@ -980,14 +980,6 @@ murfey.server.api.session_info.tomo_router:
         type: str
     methods:
       - GET
-murfey.server.api.spa.router:
-  - path: /sessions/{session_id}/cryolo_model
-    function: get_cryolo_model_path
-    path_params:
-      - name: session_id
-        type: int
-    methods:
-      - GET
 murfey.server.api.websocket.ws:
   - path: /ws/test/{client_id}
     function: websocket_endpoint

src/murfey/util/route_manifest.yaml

@@ -24,7 +24,7 @@
 from murfey.util.db import Session as MurfeySession
 from murfey.util.db import SPAFeedbackParameters, SPARelionParameters
 from murfey.util.models import FoilHoleParameters, GridSquareParameters
-from murfey.util.processing_params import default_spa_parameters
+from murfey.util.processing_params import cryolo_model_path, default_spa_parameters
 from murfey.util.spa_metadata import (
     GridSquareInfo,
     foil_hole_data,
@@ -314,15 +314,12 @@ def flush_spa_preprocess(message: dict, murfey_db: Session, demo: bool = False)
     ).all()
     if not stashed_files:
         return True
-    instrument_name = (
-        murfey_db.exec(
-            select(MurfeySession).where(MurfeySession.id == message["session_id"])
-        )
-        .one()
-        .instrument_name
-    )
-    machine_config = get_machine_config(instrument_name=instrument_name)[
-        instrument_name
+
+    murfey_session = murfey_db.exec(
+        select(MurfeySession).where(MurfeySession.id == message["session_id"])
+    ).one()
+    machine_config = get_machine_config(instrument_name=murfey_session.instrument_name)[
+        murfey_session.instrument_name
     ]
     recipe_name = machine_config.recipes.get("em-spa-preprocess", "em-spa-preprocess")
     collected_ids = murfey_db.exec(
@@ -424,6 +421,11 @@ def flush_spa_preprocess(message: dict, murfey_db: Session, demo: bool = False)
                     else f.eer_fractionation_file
                 ),
                 "do_icebreaker_jobs": default_spa_parameters.do_icebreaker_jobs,
+                "cryolo_model_weights": str(
+                    cryolo_model_path(
+                        murfey_session.visit, murfey_session.instrument_name
+                    )
+                ),
                 "foil_hole_id": foil_hole_id,
             },
         }