diff --git a/src/murfey/client/contexts/atlas.py b/src/murfey/client/contexts/atlas.py index 14e86af0f..83346d869 100644 --- a/src/murfey/client/contexts/atlas.py +++ b/src/murfey/client/contexts/atlas.py @@ -167,6 +167,30 @@ def post_transfer_epu( except KeyError: logger.info("Unable to read grid square locations from Atlas.dm") return + for p in transferred_file.parts: + if p.startswith("Sample"): + sample = int(p.replace("Sample", "")) + break + else: + logger.warning(f"Sample could not be identified for {transferred_file}") + return + + # Make sure a dcg is requested before doing grid squares + capture_post( + base_url=str(environment.url.geturl()), + router_name="workflow.router", + function_name="register_dc_group", + token=self._token, + instrument_name=environment.instrument_name, + visit_name=environment.visit, + session_id=environment.murfey_session, + data={ + "experiment_type_id": 44, # Atlas + "tag": str(transferred_file.parent), + "sample": sample, + }, + ) + # Register all grid squares on this atlas for gs, pos_data in gs_pix_positions.items(): if pos_data: capture_post( @@ -179,6 +203,7 @@ def post_transfer_epu( gsid=int(gs), data={ "tag": str(transferred_file.parent), + "sample": sample, "x_location": pos_data[0], "y_location": pos_data[1], "x_stage_position": pos_data[2], @@ -188,16 +213,8 @@ def post_transfer_epu( "angle": pos_data[6], }, ) + # Register atlas in smartem if gs_pix_positions: - for p in transferred_file.parts: - if p.startswith("Sample"): - sample = int(p.replace("Sample", "")) - break - else: - logger.warning( - f"Sample could not be identified for {transferred_file}" - ) - return capture_post( base_url=str(environment.url.geturl()), router_name="session_control.spa_router", diff --git a/src/murfey/util/models.py b/src/murfey/util/models.py index 8a8b4ed07..a13118e4d 100644 --- a/src/murfey/util/models.py +++ b/src/murfey/util/models.py @@ -272,6 +272,7 @@ class Base(BaseModel): class GridSquareParameters(BaseModel): tag: str image: str = "" + sample: int | None = None # Actual coordinates for image centre in real space x_location: Optional[float] = None diff --git a/src/murfey/workflows/spa/flush_spa_preprocess.py b/src/murfey/workflows/spa/flush_spa_preprocess.py index e90a74bc1..f3b69073e 100644 --- a/src/murfey/workflows/spa/flush_spa_preprocess.py +++ b/src/murfey/workflows/spa/flush_spa_preprocess.py @@ -65,10 +65,22 @@ def register_grid_square( if grid_square_params.width is not None: grid_square_params.width_scaled = int(grid_square_params.width / 7.8) + if grid_square_params.sample is not None: + dcg = murfey_db.exec( + select(DataCollectionGroup) + .where(DataCollectionGroup.session_id == session_id) + .where(DataCollectionGroup.sample == grid_square_params.sample) + ).one() + else: + dcg = murfey_db.exec( + select(DataCollectionGroup) + .where(DataCollectionGroup.session_id == session_id) + .where(DataCollectionGroup.tag == grid_square_params.tag) + ).one() grid_square_query = murfey_db.exec( select(GridSquare) .where(GridSquare.name == gsid) - .where(GridSquare.tag == grid_square_params.tag) + .where(GridSquare.tag == dcg.tag) .where(GridSquare.session_id == session_id) ).all() if grid_square_query: @@ -101,11 +113,6 @@ def register_grid_square( else: # No existing grid square in the murfey database if _transport_object: - dcg = murfey_db.exec( - select(DataCollectionGroup) - .where(DataCollectionGroup.session_id == session_id) - .where(DataCollectionGroup.tag == grid_square_params.tag) - ).one() gs_ispyb_response = _transport_object.do_insert_grid_square( dcg.atlas_id, gsid, grid_square_params ) @@ -149,12 +156,7 @@ def register_grid_square( instrument_name=murfey_session.instrument_name )[murfey_session.instrument_name] if machine_config.smartem_api_url: - dcg = murfey_db.exec( - select(DataCollectionGroup) - .where(DataCollectionGroup.session_id == session_id) - .where(DataCollectionGroup.tag == grid_square_params.tag) - ).one_or_none() - if dcg and dcg.smartem_grid_uuid: + if dcg.smartem_grid_uuid: secured_grid_square_image_path_full_res: Path | None = None if grid_square_params.image: secured_grid_square_image_path_full_res = secure_path( diff --git a/tests/client/contexts/test_atlas.py b/tests/client/contexts/test_atlas.py index 1989a4b4c..0ebc92558 100644 --- a/tests/client/contexts/test_atlas.py +++ b/tests/client/contexts/test_atlas.py @@ -139,7 +139,21 @@ def test_atlas_context_dm(mock_capture_post, tmp_path): context = AtlasContext("tomo", tmp_path, {}, "token") context.post_transfer(atlas_dm, environment=env) - assert mock_capture_post.call_count == 5 + assert mock_capture_post.call_count == 6 + mock_capture_post.assert_any_call( + base_url="http://localhost:8000", + router_name="workflow.router", + function_name="register_dc_group", + token="token", + instrument_name="m01", + visit_name="cm12345-6", + session_id=1, + data={ + "experiment_type_id": 44, # Atlas + "tag": str(atlas_dm.parent), + "sample": 2, + }, + ) mock_capture_post.assert_any_call( base_url="http://localhost:8000", router_name="session_control.spa_router", @@ -150,6 +164,7 @@ def test_atlas_context_dm(mock_capture_post, tmp_path): gsid=101, data={ "tag": str(atlas_dm.parent), + "sample": 2, "x_location": 1200, "y_location": 1500, "x_stage_position": 2e9, diff --git a/tests/workflows/spa/test_flush_spa_preprocess.py b/tests/workflows/spa/test_flush_spa_preprocess.py index 04c03a786..a77a785ba 100644 --- a/tests/workflows/spa/test_flush_spa_preprocess.py +++ b/tests/workflows/spa/test_flush_spa_preprocess.py @@ -14,6 +14,14 @@ def test_register_grid_square_update_add_locations( ): """Test the updating of an existing grid square""" # Create a grid square to update + dcg = DataCollectionGroup( + id=1, + session_id=ExampleVisit.murfey_session_id, + tag="session_tag", + atlas_id=90, + sample=2, + ) + murfey_db_session.add(dcg) grid_square = GridSquare( id=1, name=101, @@ -58,6 +66,14 @@ def test_register_grid_square_update_add_nothing( ): """Test the updating of an existing grid square, but with nothing to update with""" # Create a grid square to update + dcg = DataCollectionGroup( + id=1, + session_id=ExampleVisit.murfey_session_id, + tag="session_tag", + atlas_id=90, + sample=2, + ) + murfey_db_session.add(dcg) grid_square = GridSquare( id=1, name=101, @@ -71,8 +87,8 @@ def test_register_grid_square_update_add_nothing( murfey_db_session.add(grid_square) murfey_db_session.commit() - # Parameters to update with - new_parameters = GridSquareParameters(tag="session_tag") + # Parameters to update with - use a new tag but same sample + new_parameters = GridSquareParameters(tag="atlas_tag", sample=2) # Run the registration flush_spa_preprocess.register_grid_square( @@ -95,13 +111,14 @@ def test_register_grid_square_insert_with_ispyb( mock_transport, murfey_db_session: Session, tmp_path ): # Create a data collection group for lookups - grid_square = DataCollectionGroup( + dcg = DataCollectionGroup( id=1, session_id=ExampleVisit.murfey_session_id, tag="session_tag", atlas_id=90, + sample=2, ) - murfey_db_session.add(grid_square) + murfey_db_session.add(dcg) murfey_db_session.commit() # Set the ispyb return