Test the route without diameter supplied

Author: stephen-riggs

src/murfey/workflows/tomo/picking.py

@@ -114,7 +114,11 @@ def _register_picked_tomogram_use_diameter(message: dict, _db: Session):
                 class_uuids = {
                     str(i + 1): m
                     for i, m in enumerate(
-                        _murfey_id(_app_id(pj_id, _db), _db, number=50)
+                        _murfey_id(
+                            _app_id(pj_id, _db),
+                            _db,
+                            number=default_tomo_parameters.nr_classes_2d,
+                        )
                     )
                 }
                 class2d_grp_uuid = _murfey_id(_app_id(pj_id, _db), _db)[0]

tests/workflows/tomo/test_tomo_picking.py

@@ -4,6 +4,7 @@
 
 from murfey.util.db import (
     AutoProcProgram,
+    ClassificationFeedbackParameters,
     DataCollection,
     DataCollectionGroup,
     ParticleSizes,
@@ -15,15 +16,8 @@
 from tests.conftest import ExampleVisit, get_or_create_db_entry
 
 
-@mock.patch("murfey.workflows.tomo.picking._transport_object")
-@mock.patch("murfey.workflows.tomo.picking._ids_tomo_classification")
-def test_picked_tomogram_not_run_class2d(
-    mock_ids, mock_transport, murfey_db_session: Session, tmp_path
-):
-    """Run the picker feedback with less particles than needed for classification"""
-    mock_ids.return_value = [2, 1]
-
-    # Insert table dependencies
+def set_up_picking_db(murfey_db_session):
+    # Insert common elements needed in all picking tests
     dcg_entry: DataCollectionGroup = get_or_create_db_entry(
         murfey_db_session,
         DataCollectionGroup,
@@ -42,7 +36,7 @@ def test_picked_tomogram_not_run_class2d(
             "dcg_id": dcg_entry.id,
         },
     )
-    get_or_create_db_entry(
+    processing_job_entry: ProcessingJob = get_or_create_db_entry(
         murfey_db_session,
         ProcessingJob,
         lookup_kwargs={
@@ -51,6 +45,41 @@ def test_picked_tomogram_not_run_class2d(
             "dc_id": dc_entry.id,
         },
     )
+    get_or_create_db_entry(
+        murfey_db_session,
+        AutoProcProgram,
+        lookup_kwargs={
+            "id": 0,
+            "pj_id": processing_job_entry.id,
+        },
+    )
+    get_or_create_db_entry(
+        murfey_db_session,
+        ClassificationFeedbackParameters,
+        lookup_kwargs={
+            "pj_id": processing_job_entry.id,
+            "estimate_particle_diameter": True,
+            "hold_class2d": False,
+            "hold_class3d": False,
+            "class_selection_score": 0,
+            "star_combination_job": 0,
+            "initial_model": "",
+            "next_job": 0,
+        },
+    )
+    return dcg_entry.id, dc_entry.id, processing_job_entry.id
+
+
+@mock.patch("murfey.workflows.tomo.picking._transport_object")
+@mock.patch("murfey.workflows.tomo.picking._ids_tomo_classification")
+def test_picked_tomogram_not_run_class2d(
+    mock_ids, mock_transport, murfey_db_session: Session, tmp_path
+):
+    """Run the picker feedback with less particles than needed for classification"""
+    mock_ids.return_value = [2, 1]
+
+    # Insert table dependencies
+    set_up_picking_db(murfey_db_session)
 
     message = {
         "program_id": 0,
@@ -61,7 +90,6 @@ def test_picked_tomogram_not_run_class2d(
         "register": "picked_tomogram",
         "tomogram": f"{tmp_path}/Tomograms/job006/tomograms/sample.mrc",
     }
-
     picking._register_picked_tomogram_use_diameter(message, murfey_db_session)
 
     mock_ids.assert_called_once_with(0, "em-tomo-class2d", murfey_db_session)
@@ -86,53 +114,19 @@ def test_picked_tomogram_not_run_class2d(
 
 @mock.patch("murfey.workflows.tomo.picking._transport_object")
 @mock.patch("murfey.workflows.tomo.picking._ids_tomo_classification")
-def test_picked_tomogram_run_class2d(
+def test_picked_tomogram_run_class2d_with_diameter(
     mock_ids, mock_transport, murfey_db_session: Session, tmp_path
 ):
-    """Run the picker feedback with less particles than needed for classification"""
+    """Run the picker feedback with a pre-determined particle diameter"""
     mock_transport.feedback_queue = "murfey_feedback"
 
     # Insert table dependencies
-    dcg_entry: DataCollectionGroup = get_or_create_db_entry(
-        murfey_db_session,
-        DataCollectionGroup,
-        lookup_kwargs={
-            "id": 0,
-            "session_id": ExampleVisit.murfey_session_id,
-            "tag": "test_dcg",
-        },
-    )
-    dc_entry: DataCollection = get_or_create_db_entry(
-        murfey_db_session,
-        DataCollection,
-        lookup_kwargs={
-            "id": 0,
-            "tag": "test_dc",
-            "dcg_id": dcg_entry.id,
-        },
-    )
-    processing_job_entry: ProcessingJob = get_or_create_db_entry(
-        murfey_db_session,
-        ProcessingJob,
-        lookup_kwargs={
-            "id": 1,
-            "recipe": "test_recipe",
-            "dc_id": dc_entry.id,
-        },
-    )
-    get_or_create_db_entry(
-        murfey_db_session,
-        AutoProcProgram,
-        lookup_kwargs={
-            "id": 0,
-            "pj_id": processing_job_entry.id,
-        },
-    )
+    dcg_id, dc_id, pj_id = set_up_picking_db(murfey_db_session)
     get_or_create_db_entry(
         murfey_db_session,
         TomographyProcessingParameters,
         lookup_kwargs={
-            "dcg_id": dcg_entry.id,
+            "dcg_id": dcg_id,
             "pixel_size": 1.34,
             "dose_per_frame": 1,
             "frame_count": 5,
@@ -147,12 +141,12 @@ def test_picked_tomogram_run_class2d(
             ParticleSizes,
             lookup_kwargs={
                 "id": particle,
-                "pj_id": processing_job_entry.id,
+                "pj_id": pj_id,
                 "particle_size": 100,
             },
         )
 
-    mock_ids.return_value = [dcg_entry.id, 1]
+    mock_ids.return_value = [dcg_id, 1]
 
     message = {
         "session_id": 1,
@@ -164,17 +158,6 @@ def test_picked_tomogram_run_class2d(
         "register": "picked_tomogram",
         "tomogram": f"{tmp_path}/Tomograms/job006/tomograms/sample.mrc",
     }
-
-    # Create a data collection group for lookups
-    grid_square = DataCollectionGroup(
-        id=1,
-        session_id=ExampleVisit.murfey_session_id,
-        tag="session_tag",
-        atlas_id=90,
-    )
-    murfey_db_session.add(grid_square)
-    murfey_db_session.commit()
-
     picking._register_picked_tomogram_use_diameter(message, murfey_db_session)
 
     mock_ids.assert_called_once_with(0, "em-tomo-class2d", murfey_db_session)
@@ -204,6 +187,121 @@ def test_picked_tomogram_run_class2d(
                 "picker_id": None,
                 "class2d_grp_uuid": 6,
                 "class_uuids": {str(i): i for i in range(1, 6)},
+                "next_job": 9,
+                "feedback_queue": "murfey_feedback",
+            },
+            "recipes": ["em-tomo-class2d"],
+        },
+        new_connection=True,
+    )
+
+
+@mock.patch("murfey.workflows.tomo.picking._transport_object")
+@mock.patch("murfey.workflows.tomo.picking._ids_tomo_classification")
+def test_picked_tomogram_run_class2d_estimate_diameter(
+    mock_ids, mock_transport, murfey_db_session: Session, tmp_path
+):
+    """Run the picker feedback for Class2D, including diameter estimation"""
+    mock_transport.feedback_queue = "murfey_feedback"
+
+    # Insert table dependencies
+    dcg_id, dc_id, pj_id = set_up_picking_db(murfey_db_session)
+    get_or_create_db_entry(
+        murfey_db_session,
+        TomographyProcessingParameters,
+        lookup_kwargs={
+            "dcg_id": dcg_id,
+            "pixel_size": 1.34,
+            "dose_per_frame": 1,
+            "frame_count": 5,
+            "tilt_axis": 0,
+            "voltage": 300,
+            "particle_diameter": None,
+        },
+    )
+    for particle in range(10001):
+        get_or_create_db_entry(
+            murfey_db_session,
+            ParticleSizes,
+            lookup_kwargs={
+                "id": particle,
+                "pj_id": pj_id,
+                "particle_size": 100,
+            },
+        )
+    # Insert one existing tomogram which should get flushed out
+    get_or_create_db_entry(
+        murfey_db_session,
+        TomogramPicks,
+        lookup_kwargs={
+            "pj_id": pj_id,
+            "tomogram": f"{tmp_path}/Tomograms/job006/tomograms/tomogram1.mrc",
+            "cbox_3d": f"{tmp_path}/AutoPick/job007/CBOX_3d/tomogram1_picks.cbox",
+            "particle_count": 10,
+            "tomogram_pixel_size": 5.3,
+        },
+    )
+
+    mock_ids.return_value = [dcg_id, 1]
+
+    message = {
+        "session_id": 1,
+        "program_id": 0,
+        "cbox_3d": f"{tmp_path}/AutoPick/job007/CBOX_3d/sample.cbox",
+        "particle_count": 2,
+        "particle_diameters": [10.1, 20.2],
+        "pixel_size": 5.3,
+        "register": "picked_tomogram",
+        "tomogram": f"{tmp_path}/Tomograms/job006/tomograms/sample.mrc",
+    }
+    picking._register_picked_tomogram_use_diameter(message, murfey_db_session)
+
+    mock_ids.assert_called_once_with(0, "em-tomo-class2d", murfey_db_session)
+
+    # Two mock calls - one flushed tomogram and one new
+    assert mock_transport.send.call_count == 2
+    mock_transport.send.assert_any_call(
+        "processing_recipe",
+        {
+            "parameters": {
+                "tomogram": f"{tmp_path}/Tomograms/job006/tomograms/tomogram1.mrc",
+                "cbox_3d": f"{tmp_path}/AutoPick/job007/CBOX_3d/tomogram1_picks.cbox",
+                "pixel_size": 5.3,
+                "particle_diameter": 100.0,
+                "kv": 300,
+                "node_creator_queue": "node_creator",
+                "session_id": message["session_id"],
+                "autoproc_program_id": 0,
+                "batch_size": 10000,
+                "nr_classes": 5,
+                "picker_id": None,
+                "class2d_grp_uuid": 6,
+                "class_uuids": {str(i): i for i in range(1, 6)},
+                "next_job": 9,
+                "feedback_queue": "murfey_feedback",
+            },
+            "recipes": ["em-tomo-class2d"],
+        },
+        new_connection=True,
+    )
+    mock_transport.send.assert_any_call(
+        "processing_recipe",
+        {
+            "parameters": {
+                "tomogram": message["tomogram"],
+                "cbox_3d": message["cbox_3d"],
+                "pixel_size": message["pixel_size"],
+                "particle_diameter": 100.0,
+                "kv": 300,
+                "node_creator_queue": "node_creator",
+                "session_id": message["session_id"],
+                "autoproc_program_id": 0,
+                "batch_size": 10000,
+                "nr_classes": 5,
+                "picker_id": None,
+                "class2d_grp_uuid": 12,
+                "class_uuids": {str(i): i for i in range(7, 12)},
+                "next_job": 11,
                 "feedback_queue": "murfey_feedback",
             },
             "recipes": ["em-tomo-class2d"],