Move tomogram motio correction endpoint

Author: gfrn

src/pato/crud/tomograms.py

@@ -1,7 +1,7 @@
 import re
 from typing import Literal, Optional
 
-from fastapi import HTTPException, status
+from fastapi import HTTPException
 from lims_utils.tables import (
     CTF,
     MotionCorrection,
@@ -10,7 +10,7 @@
     TiltImageAlignment,
     Tomogram,
 )
-from sqlalchemy import Column, func, literal_column, select
+from sqlalchemy import Column, func, select
 from sqlalchemy import func as f
 
 from ..models.response import (
@@ -87,48 +87,46 @@ def get_shift_plot(tomogramId: int):
 
 
 def get_motion_correction(
-    limit: int, page: int, tomogramId: int, getMiddle: bool
+    limit: int, page: int, data_collection_id: int, getMiddle: bool
 ) -> FullMovieWithTilt:
-    if getMiddle:
-        total = db.session.scalar(
-            select(f.count(literal_column("1")))
-            .filter(TiltImageAlignment.tomogramId == tomogramId)
-            .order_by(TiltImageAlignment.refinedTiltAngle)
-        )
-
-        if not total:
-            raise HTTPException(
-                status_code=status.HTTP_404_NOT_FOUND,
-                detail="No items found",
-            )
-
-        limit = 1
-        page = (total // 2) - 1
-
     query = (
         select(MotionCorrection, TiltImageAlignment, CTF, Movie)
-        .filter(TiltImageAlignment.tomogramId == tomogramId)
+        .select_from(Movie)
+        .filter(Movie.dataCollectionId == data_collection_id)
         .join(
             MotionCorrection,
-            MotionCorrection.movieId == TiltImageAlignment.movieId,
+            MotionCorrection.movieId == Movie.movieId,
         )
-        .join(Movie, Movie.movieId == TiltImageAlignment.movieId)
         .join(CTF, CTF.motionCorrectionId == MotionCorrection.motionCorrectionId)
-        .order_by(TiltImageAlignment.refinedTiltAngle)
+        .outerjoin(TiltImageAlignment, TiltImageAlignment.movieId == Movie.movieId)
+        .group_by(Movie.movieId)
+        # Display movies without alignments last
+        .order_by(
+            TiltImageAlignment.refinedTiltAngle.is_(None), TiltImageAlignment.refinedTiltAngle, Movie.movieId.desc()
+        )
     )
 
-    motion = dict(db.paginate(query, limit, page))
-
     raw_total = db.session.execute(
         select(
             f.count(Movie.movieId).label("total"),
-        )
-        .select_from(Tomogram)
-        .filter(Tomogram.tomogramId == tomogramId)
-        .join(Movie, Movie.dataCollectionId == Tomogram.dataCollectionId)
+        ).filter(Movie.dataCollectionId == data_collection_id)
     ).scalar_one()
 
-    return FullMovieWithTilt(**motion, rawTotal=raw_total)
+    aligned_total = db.session.scalar(
+        select(f.count(Movie.movieId))
+        .select_from(Movie)
+        .filter(Movie.dataCollectionId == data_collection_id)
+        .join(TiltImageAlignment, TiltImageAlignment.movieId == Movie.movieId)
+    )
+
+    if getMiddle:
+        # Get central slice, ignore dark images
+        page = aligned_total // 2 - 1
+        limit = 1
+
+    motion = dict(db.paginate(query, limit, page, precounted_total=raw_total))
+
+    return FullMovieWithTilt(**motion, alignedTotal=aligned_total)
 
 
 def get_ctf(tomogramId: int):
@@ -157,6 +155,4 @@ def get_movie_path(tomogramId: int, movie_type: MovieType):
 
 
 def get_projection_path(tomogramId: int, axis: Literal["xy", "xz"]):
-    return _get_generic_tomogram_file(
-        tomogramId, Tomogram.projXZ if axis == "xz" else Tomogram.projXY
-    )
+    return _get_generic_tomogram_file(tomogramId, Tomogram.projXZ if axis == "xz" else Tomogram.projXY)

src/pato/models/response.py

@@ -194,7 +194,7 @@ class FullMovie(OrmBaseModel):
 
 
 class FullMovieWithTilt(Paged[FullMovie]):
-    rawTotal: int
+    alignedTotal: int
 
 
 class CtfBase(OrmBaseModel):

src/pato/routes/collections.py

@@ -1,11 +1,12 @@
-from fastapi import APIRouter, Body, Depends, status
+from fastapi import APIRouter, Body, Depends, Query, status
 from fastapi.responses import FileResponse
 from lims_utils.models import Paged, pagination
 
 from ..auth import Permissions
 from ..crud import collection_report as report_crud
 from ..crud import collections as crud
 from ..crud import generic
+from ..crud import tomograms as tomograms_crud
 from ..models.collections import BaseDataCollectionOut, DataCollectionFileAttachmentOut
 from ..models.reprocessing import (
     SPAReprocessingParameters,
@@ -14,6 +15,7 @@
 from ..models.response import (
     DataPoint,
     FullMovie,
+    FullMovieWithTilt,
     ItemList,
     ProcessingJobResponse,
     ReprocessingResponse,
@@ -89,6 +91,19 @@ def get_motion_correction(
     return crud.get_motion_correction(collectionId=collectionId, **page)
 
 
+@router.get("/{collectionId}/tomogram-motion", response_model=FullMovieWithTilt)
+def get_motion_correction_for_tomogram(
+    collectionId: int = Depends(auth),
+    page: dict[str, int] = Depends(pagination),
+    getMiddle=Query(
+        False,
+        description="Get index closest to the middle. Limit is set to 1, page is ignored",
+    ),
+):
+    """Get motion correction and tilt alignment data for tomogram"""
+    return tomograms_crud.get_motion_correction(data_collection_id=collectionId, getMiddle=getMiddle, **page)
+
+
 @router.get(
     "/{collectionId}/iceThickness",
 )

src/pato/routes/tomograms.py

@@ -1,15 +1,13 @@
 from typing import Literal
 
-from fastapi import APIRouter, Depends, Query
+from fastapi import APIRouter, Depends
 from fastapi.responses import FileResponse
-from lims_utils.models import pagination
 
 from ..auth import Permissions
 from ..crud import tomograms as crud
 from ..models.response import (
     CtfTiltAlign,
     DataPoint,
-    FullMovieWithTilt,
     ItemList,
 )
 from ..utils.generic import MovieType
@@ -28,21 +26,6 @@ def get_shift_plot(tomogramId: int = Depends(auth)):
     return crud.get_shift_plot(tomogramId)
 
 
-@router.get("/{tomogramId}/motion", response_model=FullMovieWithTilt)
-def get_motion_correction(
-    tomogramId: int = Depends(auth),
-    page: dict[str, int] = Depends(pagination),
-    getMiddle=Query(
-        False,
-        description="Get index closest to the middle. Limit is set to 1, page is ignored",
-    ),
-):
-    """Get motion correction data for the given tomogram"""
-    return crud.get_motion_correction(
-        tomogramId=tomogramId, getMiddle=getMiddle, **page
-    )
-
-
 @router.get("/{tomogramId}/centralSlice", response_class=FileResponse)
 def get_slice(tomogramId: int = Depends(auth), movieType: MovieType = None):
     """Get tomogram central slice image"""

tests/tomograms/test_motion.py tests/collections/test_tomogram_motion.pytests/tomograms/test_motion.py renamed to tests/collections/test_tomogram_motion.py

@@ -1,34 +1,38 @@
 def test_get(mock_permissions, client):
     """Get motion correction in a tomogram (request for admin)"""
-    resp = client.get("/tomograms/1/motion")
+    resp = client.get("/dataCollections/6017406/tomogram-motion")
     resp_json = resp.json()
 
     assert resp.status_code == 200
-    assert resp_json["total"] == 4
-    assert resp_json["rawTotal"] == 5
+    assert resp_json["total"] == 5
+    assert resp_json["alignedTotal"] == 5
 
 
 def test_get_middle(mock_permissions, client):
     """Get motion correction without tilt alignment data"""
-    resp = client.get("/tomograms/1/motion?getMiddle=true")
+    resp = client.get("/dataCollections/6017406/tomogram-motion?getMiddle=true")
     assert resp.status_code == 200
-    assert resp.json()["items"][0]["TiltImageAlignment"]["refinedTiltAngle"] == 17
+    assert resp.json()["items"][0]["TiltImageAlignment"]["refinedTiltAngle"] == 16
 
 
 def test_no_tilt_alignment(mock_permissions, client):
     """Get motion correction without tilt alignment data"""
-    resp = client.get("/tomograms/3/motion")
-    motion_correction = resp.json()
+    resp = client.get("/dataCollections/6017411/tomogram-motion")
 
     assert resp.status_code == 200
-    assert motion_correction["total"] == 0
+
+    motion_correction = resp.json()
+    assert motion_correction["total"] == 5
+    assert motion_correction["alignedTotal"] == 0
 
 
 def test_nth_motion(mock_permissions, client):
     """Get specific motion correction"""
-    resp = client.get("/tomograms/2/motion?page=1&limit=1")
-    resp_json = resp.json()
+    resp = client.get("/dataCollections/6017408/tomogram-motion?page=1&limit=1")
 
     assert resp.status_code == 200
-    assert resp_json["total"] == 4
-    assert resp_json["items"][0]["TiltImageAlignment"]["refinedTiltAngle"] == 17
+
+    motion_correction = resp.json()
+
+    assert motion_correction["total"] == 6
+    assert motion_correction["items"][0]["TiltImageAlignment"]["refinedTiltAngle"] == 18