diff --git a/pyproject.toml b/pyproject.toml index 316e3b0e..4b86a838 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -47,6 +47,7 @@ dependencies = [ "starfile", "stomp-py>8.1.1", "tifffile", # CLEM workflow + "txrm2tiff", "workflows>=3", ] [project.optional-dependencies] diff --git a/recipes/ispyb/sxt-aretomo.json b/recipes/ispyb/sxt-aretomo.json new file mode 100644 index 00000000..d4736a0c --- /dev/null +++ b/recipes/ispyb/sxt-aretomo.json @@ -0,0 +1,100 @@ +{ + "1": { + "output": 2, + "parameters": { + "dcid": "{dcid}", + "ispyb_command": "insert_tomogram", + "program_id": "{appid}", + "store_result": "ispyb_tomogram_id" + }, + "queue": "ispyb_connector", + "service": "EMISPyB" + }, + "2": { + "output": { + "denoise": 5, + "failure": 3, + "images": 6, + "ispyb_connector": 4, + "movie": 6, + "projxy": 6, + "projxz": 6, + "success": 7 + }, + "parameters": { + "dark_tol": 0, + "manual_tilt_offset": 0, + "out_bin": 1, + "pixel_size": 100, + "relion_options": {}, + "stack_file": "{stack_file}", + "tilt_axis": 0, + "txrm_file": "{txrm_file}", + "wbp": 1 + }, + "queue": "tomo_align", + "service": "TomoAlign" + }, + "3": { + "parameters": { + "ispyb_command": "update_processing_status", + "status_message": "processing failure", + "program_id": "{appid}", + "status": "failure" + }, + "queue": "ispyb_connector", + "service": "EMISPyB" + }, + "4": { + "parameters": { + "dcid": "{dcid}", + "ispyb_command": "multipart_message", + "program_id": "{appid}", + "tomogram_id": "$ispyb_tomogram_id" + }, + "queue": "ispyb_connector", + "service": "EMISPyB" + }, + "5": { + "output": { + "images": 6, + "ispyb_connector": 9, + "movie": 6, + "segmentation": 8 + }, + "queue": "denoise", + "service": "Denoise" + }, + "6": { + "queue": "images", + "service": "Images" + }, + "7": { + "parameters": { + "ispyb_command": "update_processing_status", + "status_message": "processing successful", + "program_id": "{appid}", + "status": "success" + }, + "queue": "ispyb_connector", + "service": "EMISPyB" + }, + "8": { + "output": { + "images": 6, + "ispyb_connector": 9, + "movie": 6 + }, + "queue": "segmentation", + "service": "MembrainSeg" + }, + "9": { + "parameters": { + "ispyb_command": "insert_processed_tomogram", + "tomogram_id": "$ispyb_tomogram_id" + }, + "queue": "ispyb_connector", + "service": "EMISPyB" + }, + "start": [[1, []]] +} diff --git a/src/cryoemservices/services/tomo_align.py b/src/cryoemservices/services/tomo_align.py index 45a875f3..40ccaa19 100644 --- a/src/cryoemservices/services/tomo_align.py +++ b/src/cryoemservices/services/tomo_align.py @@ -14,6 +14,11 @@ import plotly.express as px from gemmi import cif from pydantic import BaseModel, Field, ValidationError, field_validator, model_validator +from txrm2tiff.inspector import Inspector +from txrm2tiff.main import convert_and_save +from txrm2tiff.txrm import open_txrm +from txrm2tiff.txrm_functions.general import read_stream +from txrm2tiff.xradia_properties.enums import XrmDataTypes from workflows.recipe import wrap_subscribe from cryoemservices.services.common_service import CommonService @@ -100,11 +105,12 @@ class TomoParameters(BaseModel): aretomo_version: Literal[2, 3] = 3 stack_file: str = Field(..., min_length=1) pixel_size: float + txrm_file: Optional[str] = None path_pattern: Optional[str] = None input_file_list: Optional[str] = None vol_z: Optional[int] = None max_vol: int = 1800 - min_vol: int = 800 + min_vol: int = 600 extra_vol: int = 400 out_bin: int = 4 second_bin: Optional[int] = None @@ -139,12 +145,15 @@ class TomoParameters(BaseModel): def check_only_one_is_provided(cls, values): input_file_list = values.get("input_file_list") path_pattern = values.get("path_pattern") - if not input_file_list and not path_pattern: - raise ValueError("input_file_list or path_pattern must be provided") - if input_file_list and path_pattern: + txrm_file = values.get("txrm_file") + if not input_file_list and not path_pattern and not txrm_file: raise ValueError( - "Message must only include one of 'path_pattern' and 'input_file_list'." - " Both are set or one has been set by the recipe." + "input_file_list or path_pattern or txrm_file must be provided" + ) + if input_file_list and path_pattern and txrm_file: + raise ValueError( + "Message must only include one of " + "'path_pattern', 'input_file_list' or 'txrm_file'." ) return values @@ -180,6 +189,7 @@ class TomoAlign(CommonService): # Values to extract for ISPyB input_file_list_of_lists: List[Any] + tilt_angles: dict refined_tilts: List[float] x_shift: List[float] y_shift: List[float] @@ -192,6 +202,7 @@ class TomoAlign(CommonService): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) + self.tilt_angles = {} self.refined_tilts = [] self.rot_centre_z_list = [] @@ -307,15 +318,13 @@ def _tilt(file_list_for_tilts): # If no volume provided, set it to the maximum we allow tomo_params.vol_z = tomo_params.max_vol - # Update the relion options - tomo_params.relion_options = update_relion_options( - tomo_params.relion_options, dict(tomo_params) - ) - tomo_params.relion_options.pixel_size_downscaled = ( - tomo_params.pixel_size * tomo_params.out_bin - ) + # Get the names of the output files expected + alignment_output_dir = Path(tomo_params.stack_file).parent + Path(tomo_params.stack_file).parent.mkdir(parents=True, exist_ok=True) + stack_name = str(Path(tomo_params.stack_file).stem) # Convert a path pattern into a file list + self.tilt_angles = {} if tomo_params.path_pattern: directory = Path(tomo_params.path_pattern).parent @@ -329,10 +338,34 @@ def _tilt(file_list_for_tilts): elif tomo_params.input_file_list: file_list = ast.literal_eval(tomo_params.input_file_list) self.input_file_list_of_lists = file_list + elif tomo_params.txrm_file and Path(tomo_params.txrm_file).is_file(): + tomo_params.pixel_size = self.convert_txrm_to_stack( + tomo_params.txrm_file, tomo_params.stack_file + ) + self.input_file_list_of_lists = [] + # Check we now have the expected stack file + if not Path(tomo_params.stack_file).is_file(): + self.log.warning("Stack file generation failed") + rw.transport.nack(header) + return + else: + self.log.warning(f"Invalid input or {tomo_params.txrm_file} is not a file") + rw.transport.nack(header) + return - self.log.info(f"Input list {self.input_file_list_of_lists}") + self.log.info( + f"Input list {self.input_file_list_of_lists or tomo_params.txrm_file}" + ) self.input_file_list_of_lists.sort(key=_tilt) + # Update the relion options + tomo_params.relion_options = update_relion_options( + tomo_params.relion_options, dict(tomo_params) + ) + tomo_params.relion_options.pixel_size_downscaled = ( + tomo_params.pixel_size * tomo_params.out_bin + ) + # Find all the tilt angles and remove duplicates tilt_dict: dict = {} for tilt in self.input_file_list_of_lists: @@ -386,11 +419,18 @@ def _tilt(file_list_for_tilts): self.input_file_list_of_lists.remove(self.input_file_list_of_lists[index]) # Find the input image dimensions - with mrcfile.open(self.input_file_list_of_lists[0][0]) as mrc: + if self.input_file_list_of_lists: + input_image = self.input_file_list_of_lists[0][0] + else: + input_image = tomo_params.stack_file + # Read the first image if a list is given or stack if using txrm + with mrcfile.open(input_image) as mrc: mrc_header = mrc.header # x and y get flipped on tomogram creation tomo_params.relion_options.tomo_size_x = int(mrc_header.nx) tomo_params.relion_options.tomo_size_y = int(mrc_header.ny) + + # Scale size of output scaled_x_size = tomo_params.relion_options.tomo_size_x / int( tomo_params.out_bin ) @@ -398,52 +438,48 @@ def _tilt(file_list_for_tilts): tomo_params.out_bin ) - # Get the names of the output files expected - alignment_output_dir = Path(tomo_params.stack_file).parent - Path(tomo_params.stack_file).parent.mkdir(parents=True, exist_ok=True) - stack_name = str(Path(tomo_params.stack_file).stem) - - project_dir_search = re.search(".+/job[0-9]+/", tomo_params.stack_file) + project_dir_search = re.search(".+/Tomograms/", tomo_params.stack_file) job_num_search = re.search("/job[0-9]+", tomo_params.stack_file) if project_dir_search and job_num_search: - project_dir = Path(project_dir_search[0]).parent.parent + project_dir = Path(project_dir_search[0]).parent job_number = int(job_num_search[0][4:]) + elif project_dir_search: + # Allow processing without job numbers, but then skip node creator sends + project_dir = Path(project_dir_search[0]).parent + job_number = 0 else: self.log.warning(f"Invalid project directory in {tomo_params.stack_file}") rw.transport.nack(header) return - # Stack the tilts - try: - create_tilt_stack( - self.input_file_list_of_lists, Path(tomo_params.stack_file) - ) - except (FileNotFoundError, ValueError) as e: - self.log.error(f"Creating stack file failed: {e}") - rw.transport.nack(header) - return + if self.input_file_list_of_lists: + # Stack the tilts + try: + create_tilt_stack( + self.input_file_list_of_lists, Path(tomo_params.stack_file) + ) + except (FileNotFoundError, ValueError) as e: + self.log.error(f"Creating stack file failed: {e}") + rw.transport.nack(header) + return # Set up the angle file needed for dose weighting - angle_file = ( - Path(tomo_params.stack_file).parent - / f"{Path(tomo_params.stack_file).stem}_TLT.txt" - ) - tilt_angles = {} + angle_file = Path(tomo_params.stack_file).parent / f"{stack_name}_TLT.txt" for i in range(len(self.input_file_list_of_lists)): tilt_index = get_tilt_number_v5_12( Path(self.input_file_list_of_lists[i][0]) ) tilt_index -= len(np.where(removed_tilt_numbers < tilt_index)[0]) - tilt_angles[tilt_index] = float(self.input_file_list_of_lists[i][1]) + self.tilt_angles[tilt_index] = float(self.input_file_list_of_lists[i][1]) with open(angle_file, "w") as angfile: - for tilt_id in tilt_angles.keys(): + for tilt_id in self.tilt_angles.keys(): if tomo_params.aretomo_version == 3 and tomo_params.manual_tilt_offset: # AreTomo3 performs better with pre-shifted angles angfile.write( - f"{tilt_angles[tilt_id] + tomo_params.manual_tilt_offset:.2f} {int(tilt_id)}\n" + f"{self.tilt_angles[tilt_id] + tomo_params.manual_tilt_offset:.2f} {int(tilt_id)}\n" ) else: - angfile.write(f"{tilt_angles[tilt_id]:.2f} {int(tilt_id)}\n") + angfile.write(f"{self.tilt_angles[tilt_id]:.2f} {int(tilt_id)}\n") # Do alignment with AreTomo aretomo_output_path = alignment_output_dir / f"{stack_name}_Vol.mrc" @@ -461,7 +497,6 @@ def _tilt(file_list_for_tilts): plot_path = alignment_output_dir / plot_file # Extract results - pixel_spacing: str = str(tomo_params.pixel_size * tomo_params.out_bin) aln_file = self.extract_from_aln(tomo_params, alignment_output_dir, plot_path) if not aretomo_result.returncode and not aln_file: self.log.error("Failed to read alignment file") @@ -485,13 +520,15 @@ def _tilt(file_list_for_tilts): tomo_params.vol_z = tomo_params.min_vol tomo_params.relion_options.vol_z = tomo_params.vol_z - if not job_is_rerun: + if job_number and not job_is_rerun: # Send to node creator if this is the first time this tomogram is made self.log.info("Sending tomo align to node creator") node_creator_parameters = { "experiment_type": "tomography", "job_type": self.job_type, - "input_file": f"{project_dir}/AlignTiltSeries/job{job_number - 1:03}/tilts/{Path(self.input_file_list_of_lists[0][0]).name}", + "input_file": f"{project_dir}/AlignTiltSeries/job{job_number - 1:03}/tilts/{Path(self.input_file_list_of_lists[0][0]).name}" + if self.input_file_list_of_lists + else tomo_params.txrm_file, "output_file": str(aretomo_output_path), "relion_options": dict(tomo_params.relion_options), "command": " ".join(aretomo_command), @@ -574,7 +611,7 @@ def _tilt(file_list_for_tilts): "size_x": scaled_x_size, # volume image size, pix "size_y": scaled_y_size, "size_z": scaled_z_size, - "pixel_spacing": pixel_spacing, + "pixel_spacing": str(tomo_params.relion_options.pixel_size_downscaled), "tilt_angle_offset": str( self.tilt_offset or tomo_params.manual_tilt_offset ), @@ -614,18 +651,19 @@ def _tilt(file_list_for_tilts): im_diff = 0 # TiltImageAlignment (one per movie) node_creator_params_list = [] - (project_dir / f"ExcludeTiltImages/job{job_number - 2:03}").mkdir( - parents=True, exist_ok=True - ) - if not ( - project_dir / f"ExcludeTiltImages/job{job_number - 2:03}/tilts" - ).is_symlink(): - ( + if self.input_file_list_of_lists and job_number: + (project_dir / f"ExcludeTiltImages/job{job_number - 2:03}").mkdir( + parents=True, exist_ok=True + ) + if not ( project_dir / f"ExcludeTiltImages/job{job_number - 2:03}/tilts" - ).symlink_to(project_dir / "MotionCorr/job002/Movies") - (project_dir / f"AlignTiltSeries/job{job_number - 1:03}").mkdir( - parents=True, exist_ok=True - ) + ).is_symlink(): + ( + project_dir / f"ExcludeTiltImages/job{job_number - 2:03}/tilts" + ).symlink_to(project_dir / "MotionCorr/job002/Movies") + (project_dir / f"AlignTiltSeries/job{job_number - 1:03}").mkdir( + parents=True, exist_ok=True + ) if self.rot: tomo_params.relion_options.tilt_axis_angle = self.rot for im, movie in enumerate(self.input_file_list_of_lists): @@ -701,8 +739,9 @@ def _tilt(file_list_for_tilts): rw.transport.nack(header) return - for tilt_params in node_creator_params_list: - rw.send_to("node_creator", tilt_params) + if job_number: + for tilt_params in node_creator_params_list: + rw.send_to("node_creator", tilt_params) ispyb_command_list.append( { @@ -771,7 +810,7 @@ def _tilt(file_list_for_tilts): "image_command": "mrc_projection", "file": str(aretomo_output_path), "projection": projection_type, - "pixel_spacing": float(pixel_spacing), + "pixel_spacing": tomo_params.relion_options.pixel_size_downscaled, } if projection_type in ["XZ", "YZ"] and self.thickness_pixels: images_call_params["thickness_ang"] = ( @@ -781,13 +820,20 @@ def _tilt(file_list_for_tilts): # Forward results to denoise service self.log.info(f"Sending to denoise service {aretomo_output_path}") + if job_number: + denoise_dir = project_dir / f"Denoise/job{job_number + 1:03}/tomograms" + else: + denoise_dir = project_dir / "Denoise" + if tomo_params.relion_options.pixel_size_downscaled > 20: + # This is used for SXT + # Resets the pixel size to prevent segmentation rescaling + tomo_params.relion_options.pixel_size = 10 + tomo_params.relion_options.pixel_size_downscaled = 10 rw.send_to( "denoise", { "volume": str(aretomo_output_path), - "output_dir": str( - project_dir / f"Denoise/job{job_number + 1:03}/tomograms" - ), + "output_dir": str(denoise_dir), "relion_options": dict(tomo_params.relion_options), }, ) @@ -811,6 +857,36 @@ def _tilt(file_list_for_tilts): self.log.info(f"Done tomogram alignment for {tomo_params.stack_file}") rw.transport.ack(header) + def convert_txrm_to_stack(self, txrm_file: str, stack_file: str) -> float: + # Read the tilt angles and pixel size from the txrm + with open_txrm( + txrm_file, load_images=False, load_reference=False, strict=False + ) as txrm: + inspector = Inspector(txrm) + angles = read_stream( + inspector.txrm.ole, + "ImageInfo/Angles", + XrmDataTypes.XRM_FLOAT, + strict=True, + ) + pixel_size_microns = read_stream( + inspector.txrm.ole, + "ImageInfo/PixelSize", + XrmDataTypes.XRM_FLOAT, + strict=True, + ) + self.tilt_angles = dict(enumerate(angles)) + + # Convert the txrm to a tiff stack, then convert to mrc for aretomo + tifftomo = Path(stack_file).with_suffix(".tiff") + convert_and_save(txrm_file, str(tifftomo), custom_reference=None) + # Let this run, and check later if the output file exists + subprocess.run(["tif2mrc", str(tifftomo), stack_file]) + tifftomo.unlink(missing_ok=True) + if pixel_size_microns: + return pixel_size_microns[0] * 1000 + return 100 + def assemble_aretomo3_command( self, aretomo_executable: str, diff --git a/tests/services/test_tomo_align.py b/tests/services/test_tomo_align.py index 5d8a52e7..bc6536a7 100644 --- a/tests/services/test_tomo_align.py +++ b/tests/services/test_tomo_align.py @@ -9,6 +9,7 @@ import mrcfile import numpy as np import pytest +from txrm2tiff.xradia_properties.enums import XrmDataTypes from workflows.transport.offline_transport import OfflineTransport from cryoemservices.services import tomo_align @@ -16,9 +17,10 @@ class MrcFileHeader: - def __init__(self, nx, ny): + def __init__(self, nx, ny, nz=1): self.nx = nx self.ny = ny + self.nz = nz @pytest.fixture @@ -1410,7 +1412,7 @@ def test_tomo_align_service_path_pattern( output_relion_options["frame_count"] = 6 output_relion_options["dose_per_frame"] = 0.2 # Volume default to minimum - output_relion_options["vol_z"] = 800 + output_relion_options["vol_z"] = 600 # Set up the mock service service = tomo_align.TomoAlign( @@ -1510,10 +1512,10 @@ def write_aretomo_outputs(command, capture_output): # Check resizing and rotating assert mock_resize.call_count == 2 mock_resize.assert_any_call( - tmp_path / "Tomograms/job006/tomograms/test_stack_Vol.mrc", int(800 / 4) + tmp_path / "Tomograms/job006/tomograms/test_stack_Vol.mrc", int(600 / 4) ) mock_resize.assert_any_call( - tmp_path / "Tomograms/job006/tomograms/test_stack_2ND_Vol.mrc", int(800 / 2) + tmp_path / "Tomograms/job006/tomograms/test_stack_2ND_Vol.mrc", int(600 / 2) ) mock_rotate.assert_not_called() @@ -2079,6 +2081,288 @@ def write_aretomo_outputs(command, capture_output): offline_transport.send.assert_any_call("failure", {}) +@mock.patch("cryoemservices.services.tomo_align.subprocess.run") +@mock.patch("cryoemservices.services.tomo_align.mrcfile") +@mock.patch("cryoemservices.services.tomo_align.convert_and_save") +@mock.patch("cryoemservices.services.tomo_align.open_txrm") +@mock.patch("cryoemservices.services.tomo_align.Inspector") +@mock.patch("cryoemservices.services.tomo_align.read_stream") +@mock.patch("cryoemservices.services.tomo_align.resize_tomogram") +@mock.patch("cryoemservices.services.tomo_align.rotate_tomogram") +def test_tomo_align_service_txrm( + mock_rotate, + mock_resize, + mock_read_stream, + mock_inspector, + mock_open_txrm, + mock_convert_and_save, + mock_mrcfile, + mock_subprocess, + offline_transport, + tmp_path, +): + """ + Send a test message to TomoAlign (AreTomo3) for a txrm file + """ + mock_mrcfile.open().__enter__().header = MrcFileHeader(nx=4000, ny=3000, nz=600) + mock_read_stream.return_value = [0.1, 0.3, 0.5] + + header = { + "message-id": mock.sentinel, + "subscription": mock.sentinel, + } + tomo_align_test_message = { + "stack_file": f"{tmp_path}/Tomograms/stack.mrc", + "txrm_file": f"{tmp_path}/tilt_stack.txrm", + "pixel_size": 100, + "relion_options": {}, + "dark_tol": 0, + "manual_tilt_offset": 0, + "out_bin": 1, + "tilt_axis": 0, + "wbp": 1, + } + output_relion_options = dict(RelionServiceOptions()) + output_relion_options["pixel_size"] = 100 + output_relion_options["pixel_size_downscaled"] = 100 + output_relion_options["tomo_size_x"] = 4000 + output_relion_options["tomo_size_y"] = 3000 + output_relion_options["vol_z"] = 600 + + # Touch the expected input files + (tmp_path / "tilt_stack.txrm").touch() + + # Set up the mock service + service = tomo_align.TomoAlign( + environment={"queue": ""}, transport=offline_transport + ) + service.initializing() + + def write_aretomo_outputs(command, capture_output: bool = False): + if command[0] != "AreTomo3": + (tmp_path / "Tomograms/stack.mrc").touch(exist_ok=True) + return CompletedProcess("", returncode=0) + # Set up outputs: stack_Imod file like AreTomo3, no exclusions but with space + (tmp_path / "Tomograms/stack_Imod").mkdir(parents=True) + (tmp_path / "Tomograms/stack_Vol.mrc").touch() + with open(tmp_path / "Tomograms/stack_Imod/tilt.com", "w") as dark_file: + dark_file.write("EXCLUDELIST ") + with open(tmp_path / "Tomograms/stack.aln", "w") as aln_file: + aln_file.write("# Thickness = 130\ndummy 0.0 1000 1.2 2.3 5 6 7 8 4.5") + return CompletedProcess( + "", + returncode=0, + stdout=( + "Rot center Z 100.0 200.0 3.1\n" + "Rot center Z 150.0 250.0 2.1\n" + "Tilt offset 1.1, CC: 0.5\n" + "Best tilt axis: 57, Score: 0.5\n" + ).encode("ascii"), + stderr="stderr".encode("ascii"), + ) + + mock_subprocess.side_effect = write_aretomo_outputs + + # Send a message to the service + service.tomo_align(None, header=header, message=tomo_align_test_message) + + aretomo_command = [ + "AreTomo3", + "-Cmd", + "1", + "-InPrefix", + tomo_align_test_message["stack_file"], + "-OutDir", + f"{tmp_path}/Tomograms", + "-TiltCor", + "1", + "0.0", + "-TiltAxis", + "0.0", + "1", + "-AtBin", + "1", + "-PixSize", + "100.0", + "-VolZ", + "1800", + "-ExtZ", + "400", + "-FlipVol", + "0", + "-Wbp", + "1", + "-OutImod", + "1", + "-DarkTol", + "0.0", + ] + + # Check txrm file reading + mock_open_txrm.assert_called_once_with( + tomo_align_test_message["txrm_file"], + load_images=False, + load_reference=False, + strict=False, + ) + mock_open_txrm().__enter__.assert_called_once() + mock_inspector.assert_called_once() + mock_read_stream.assert_any_call( + mock.ANY, + "ImageInfo/Angles", + XrmDataTypes.XRM_FLOAT, + strict=True, + ) + mock_read_stream.assert_any_call( + mock.ANY, + "ImageInfo/PixelSize", + XrmDataTypes.XRM_FLOAT, + strict=True, + ) + mock_convert_and_save.assert_called_once_with( + tomo_align_test_message["txrm_file"], + f"{tmp_path}/Tomograms/stack.tiff", + custom_reference=None, + ) + + # Check the expected calls were made + assert mock_subprocess.call_count == 2 + mock_subprocess.assert_any_call( + [ + "tif2mrc", + f"{tmp_path}/Tomograms/stack.tiff", + f"{tmp_path}/Tomograms/stack.mrc", + ] + ) + mock_subprocess.assert_any_call( + aretomo_command, + capture_output=True, + ) + + # Check resizing and rotating + mock_resize.assert_called_once_with(tmp_path / "Tomograms/stack_Vol.mrc", 600) + mock_rotate.assert_called_once_with(tmp_path / "Tomograms/stack_Vol.mrc", 0) + + # Check the angle file + assert (tmp_path / "Tomograms/stack_TLT.txt").is_file() + with open(tmp_path / "Tomograms/stack_TLT.txt", "r") as angfile: + angles_data = angfile.read() + assert angles_data == "0.10 0\n0.30 1\n0.50 2\n" + + # Check the shift plot + with open(tmp_path / "Tomograms/stack_xy_shift_plot.json") as shift_plot: + shift_data = json.load(shift_plot) + assert shift_data["data"][0]["x"] == [1.2] + assert shift_data["data"][0]["y"] == [2.3] + + # Check that the correct messages were sent + assert offline_transport.send.call_count == 10 + offline_transport.send.assert_any_call( + "ispyb_connector", + { + "ispyb_command": "multipart_message", + "ispyb_command_list": [ + { + "ispyb_command": "insert_tomogram", + "volume_file": "stack_Vol.mrc", + "stack_file": tomo_align_test_message["stack_file"], + "size_x": 4000.0, + "size_y": 3000.0, + "size_z": 600, + "pixel_spacing": "100.0", + "tilt_angle_offset": "1.1", + "z_shift": "2.1", + "file_directory": f"{tmp_path}/Tomograms", + "central_slice_image": "stack_Vol_thumbnail.jpeg", + "tomogram_movie": "stack_Vol_movie.png", + "xy_shift_plot": "stack_xy_shift_plot.json", + "proj_xy": "stack_Vol_projXY.jpeg", + "proj_xz": "stack_Vol_projXZ.jpeg", + "alignment_quality": "0.5", + }, + { + "ispyb_command": "insert_processed_tomogram", + "file_path": f"{tmp_path}/Tomograms/stack.mrc", + "processing_type": "Stack", + }, + { + "ispyb_command": "insert_processed_tomogram", + "file_path": f"{tmp_path}/Tomograms/stack_alignment.jpeg", + "processing_type": "Alignment", + }, + ], + }, + ) + offline_transport.send.assert_any_call( + "images", + { + "image_command": "tilt_series_alignment", + "file": tomo_align_test_message["stack_file"], + "aln_file": f"{tmp_path}/Tomograms/stack.aln", + "pixel_size": tomo_align_test_message["pixel_size"], + }, + ) + offline_transport.send.assert_any_call( + "images", + { + "image_command": "mrc_central_slice", + "file": tomo_align_test_message["stack_file"], + }, + ) + offline_transport.send.assert_any_call( + "images", + { + "image_command": "mrc_to_apng", + "file": tomo_align_test_message["stack_file"], + }, + ) + offline_transport.send.assert_any_call( + "images", + { + "image_command": "mrc_central_slice", + "file": f"{tmp_path}/Tomograms/stack_Vol.mrc", + }, + ) + offline_transport.send.assert_any_call( + "images", + { + "image_command": "mrc_to_apng", + "file": f"{tmp_path}/Tomograms/stack_Vol.mrc", + }, + ) + offline_transport.send.assert_any_call( + "images", + { + "image_command": "mrc_projection", + "file": f"{tmp_path}/Tomograms/stack_Vol.mrc", + "projection": "XY", + "pixel_spacing": 100.0, + }, + ) + offline_transport.send.assert_any_call( + "images", + { + "image_command": "mrc_projection", + "file": f"{tmp_path}/Tomograms/stack_Vol.mrc", + "projection": "YZ", + "pixel_spacing": 100.0, + "thickness_ang": 13000.0, + }, + ) + # Denoise gets sent with reset pixel size + output_relion_options["pixel_size"] = 10 + output_relion_options["pixel_size_downscaled"] = 10 + offline_transport.send.assert_any_call( + "denoise", + { + "volume": f"{tmp_path}/Tomograms/stack_Vol.mrc", + "output_dir": f"{tmp_path}/Denoise", + "relion_options": output_relion_options, + }, + ) + offline_transport.send.assert_any_call("success", {}) + + def test_parse_tomo_align_output(offline_transport): """ Send test lines to the output parser diff --git a/tests/services/test_tomo_align_slurm.py b/tests/services/test_tomo_align_slurm.py index 11330320..68485bb9 100644 --- a/tests/services/test_tomo_align_slurm.py +++ b/tests/services/test_tomo_align_slurm.py @@ -253,7 +253,7 @@ def test_tomo_align_slurm_service_aretomo3( # Check resizing and rotating at min vol mock_resize.assert_called_once_with( tmp_path / "cm12345-6/Tomograms/job006/tomograms/test_stack_Vol.mrc", - int(800 / 4), + int(600 / 4), ) mock_rotate.assert_not_called()