diff --git a/dynamight/data/dataloaders/relion.py b/dynamight/data/dataloaders/relion.py index 1a1126d..0fc2b08 100644 --- a/dynamight/data/dataloaders/relion.py +++ b/dynamight/data/dataloaders/relion.py @@ -18,10 +18,9 @@ def is_relion_abort(directory: str) -> bool: return os.path.isfile(os.path.join(directory, "RELION_JOB_ABORT_NOW")) - def write_relion_job_exit_status( - directory: str, status: str, pipeline_control: bool = False, + directory: str, status: str, pipeline_control: bool = True, ): if pipeline_control: @@ -29,6 +28,7 @@ def write_relion_job_exit_status( elif status == "FAILURE": sys.exit(1) + def abort_if_relion_abort(directory: str): if is_relion_abort(directory): write_relion_job_exit_status(directory, "ABORTED") @@ -82,21 +82,24 @@ def load(self, path: str) -> None: data_star_path = path root_search_path = os.path.dirname(os.path.abspath(path)) else: - data_star_path = os.path.abspath(self._find_star_file_in_path(path, "data")) + data_star_path = os.path.abspath( + self._find_star_file_in_path(path, "data")) root_search_path = os.path.abspath(path) self.data_star_path = os.path.abspath(data_star_path) data = load_star(self.data_star_path) if 'optics' not in data: - raise RuntimeError("Optics groups table not found in data star file") + raise RuntimeError( + "Optics groups table not found in data star file") if 'particles' not in data: raise RuntimeError("Particles table not found in data star file") self._load_optics_group(data['optics']) self._load_particles(data['particles']) - self.project_root = self._find_project_root(root_search_path, self.image_file_paths[0]) + self.project_root = self._find_project_root( + root_search_path, self.image_file_paths[0]) # Convert image paths to absolute paths for i in range(len(self.image_file_paths)): @@ -214,7 +217,7 @@ def _load_particles(self, particles: dict) -> None: # CTF parameters ------------------------------------- if 'rlnDefocusU' in particles and \ 'rlnDefocusV' in particles and \ - 'rlnDefocusAngle' in particles: + 'rlnDefocusAngle' in particles: ctf_u = float(particles['rlnDefocusU'][i]) ctf_v = float(particles['rlnDefocusV'][i]) ctf_a = float(particles['rlnDefocusAngle'][i]) @@ -225,7 +228,7 @@ def _load_particles(self, particles: dict) -> None: # Rotation parameters -------------------------------- if 'rlnAngleRot' in particles and \ 'rlnAngleTilt' in particles and \ - 'rlnAnglePsi' in particles: + 'rlnAnglePsi' in particles: a = np.array([ float(particles['rlnAngleRot'][i]), float(particles['rlnAngleTilt'][i]), @@ -242,8 +245,10 @@ def _load_particles(self, particles: dict) -> None: # Translation parameters ------------------------------ if 'rlnOriginXAngst' in particles and 'rlnOriginYAngst' in particles: - trans_x = float(particles['rlnOriginXAngst'][i]) / og['pixel_size'] - trans_y = float(particles['rlnOriginYAngst'][i]) / og['pixel_size'] + trans_x = float( + particles['rlnOriginXAngst'][i]) / og['pixel_size'] + trans_y = float( + particles['rlnOriginYAngst'][i]) / og['pixel_size'] else: trans_x = 0. trans_y = 0. @@ -259,7 +264,8 @@ def _load_particles(self, particles: dict) -> None: image_stack_id = 0 img_path = img_tokens[1] else: - raise RuntimeError(f"Invalid image file name (rlnImageName): {img_name}") + raise RuntimeError( + f"Invalid image file name (rlnImageName): {img_name}") self.part_stack_idx.append(image_stack_id) @@ -274,7 +280,8 @@ def _load_particles(self, particles: dict) -> None: self.part_og_idx = np.array(self.part_og_idx) self.part_defocus = np.array(self.part_defocus, dtype=np.float32) self.part_rotation = np.array(self.part_rotation, dtype=np.float32) - self.part_translation = np.array(self.part_translation, dtype=np.float32) + self.part_translation = np.array( + self.part_translation, dtype=np.float32) self.part_noise_group_id = np.array(self.part_noise_group_id) self.part_stack_idx = np.array(self.part_stack_idx) self.part_image_file_path_idx = np.array(self.part_image_file_path_idx) @@ -289,7 +296,8 @@ def _find_star_file_in_path(path: str, type: str = "optimiser") -> str: files = list.sort(files) return files[-1] - raise FileNotFoundError(f"Could not find '{type}' star-file in path: {path}") + raise FileNotFoundError( + f"Could not find '{type}' star-file in path: {path}") @staticmethod def _find_project_root(from_path: str, file_relative_path: str) -> str: @@ -306,7 +314,8 @@ def _find_project_root(from_path: str, file_relative_path: str) -> str: trial_path = os.path.join(current_path, file_relative_path) if os.path.isfile(trial_path): return current_path - if current_path == os.path.dirname(current_path): # At filesystem root + # At filesystem root + if current_path == os.path.dirname(current_path): raise RuntimeError( f"Relion project directory could not be found from the subdirectory: {from_path}") current_path = os.path.dirname(current_path) @@ -318,8 +327,9 @@ def _infer_particle_diameter(self): if self.particle_diameter is None: self.particle_diameter = self.box_size * 1 * pixel_size - \ - self.circular_mask_thickness - print(f"Assigning a diameter of {round(self.particle_diameter)} angstrom") + self.circular_mask_thickness + print( + f"Assigning a diameter of {round(self.particle_diameter)} angstrom") else: if self.particle_diameter > max_diameter_ang: print(