pycistem 0.6.1__cp310-cp310-manylinux_2_28_x86_64.whl

pycistem/programs/unblur.py

@@ -0,0 +1,247 @@
+import asyncio
+import datetime
+import sqlite3
+import struct
+from dataclasses import dataclass
+from pathlib import Path
+from typing import List, Union
+import time
+
+import mrcfile
+import pandas as pd
+
+from pycistem.database import datetime_to_msdos, get_movie_info_from_db
+from pycistem.programs import cistem_program
+from pycistem.programs._cistem_constants import socket_job_result, socket_send_next_job
+
+
+@dataclass
+class UnblurParameters:
+    input_filename: str
+    output_filename: str = "unblurred.mrc"
+    pixel_size: float = 1.0
+    minimum_shift_in_angstroms: float = 2.00
+    maximum_shift_in_angstroms: float = 40.0
+    should_dose_filter: bool = True
+    should_restore_power: bool = True
+    termination_threshold_in_angstroms: float = 1.0
+    max_iterations: int = 10
+    bfactor_in_angstroms: float = 1500
+    should_mask_central_cross: bool = True
+    horizontal_mask_size: int = 1
+    vertical_mask_size: int = 1
+    acceleration_voltage: float = 300.0
+    exposure_per_frame: float = 0.0
+    pre_exposure_amount: float = 0.0
+    movie_is_gain_corrected: bool = False
+    gain_filename: str = "gain.mrc"
+    movie_is_dark_corrected: bool = True
+    dark_filename: str = "dark.mrc"
+    output_binning_factor: float = 2.0
+    correct_mag_distortion: bool = False
+    mag_distortion_angle: float = 0.0
+    mag_distortion_major_scale: float = 1.0
+    mag_distortion_minor_scale: float = 1.0
+    write_out_amplitude_spectrum: bool = True
+    amplitude_spectrum_filename: str = "amplitude_spectrum.mrc"
+    write_out_small_sum_image: bool = True
+    small_sum_image_filename: str = "scaled_sum.mrc"
+    first_frame: int = 1
+    last_frame: int = 0
+    number_of_frames_for_running_average: int = 1
+    max_threads: int = 1
+    save_aligned_frames: bool = False
+    aligned_frames_filename: str = "aligned_frames.mrc"
+    output_shift_text_file: str = "shifts.txt"
+    eer_frames_per_image: int = 0
+    eer_super_res_factor: int = 1
+    align_on_cropped_area: bool = False
+    cropped_area_center_x : int = 0
+    cropped_area_center_y : int = 0
+    cropped_area_size_x : int = -1
+    cropped_area_size_y  : int = -1
+    replace_dark_areas_with_gaussian_noise: bool = False
+    threshold_for_gaussian_noise: float = 0.1
+    measure_mean_and_variance_for_gaussian_noise : bool = False
+    mean_for_gaussian_noise : float = 0.0
+    variance_for_gaussian_noise : float = 0.0
+
+def parameters_from_database(database, decolace=False, **kwargs):
+    movie_info = get_movie_info_from_db(database)
+    ProjectDirectory = Path(database).parent
+    par = [UnblurParameters(
+        input_filename = movie["FILENAME"],
+        output_filename= (ProjectDirectory / "Assets" / "Images" / f"{Path(movie['FILENAME']).stem}_{movie['MOVIE_ASSET_ID']}_auto_{i}.mrc").as_posix(),
+        pixel_size = movie["PIXEL_SIZE"],
+        gain_filename=movie["GAIN_FILENAME"],
+        output_binning_factor=movie["OUTPUT_BINNING_FACTOR"],
+        exposure_per_frame=movie["DOSE_PER_FRAME"],
+        amplitude_spectrum_filename=(ProjectDirectory / "Assets" / "Images" / "Spectra" / f"{Path(movie['FILENAME']).stem}_{movie['MOVIE_ASSET_ID']}_auto.mrc").as_posix(),
+        small_sum_image_filename=(ProjectDirectory / "Assets" / "Images" / "Scaled" / f"{Path(movie['FILENAME']).stem}_{movie['MOVIE_ASSET_ID']}_auto.mrc").as_posix(),
+        align_on_cropped_area=decolace,
+        replace_dark_areas_with_gaussian_noise=decolace
+    ) for i,movie in movie_info.iterrows()]
+    return(par)
+
+def write_results_to_database(database,  parameters, results, change_image_assets=True):
+    conn = sqlite3.connect(database)
+    cur = conn.cursor()
+    results = sorted(results, key=lambda x: x["parameter_index"])
+    MOVIE_ALIGNMENT_LIST = []
+
+    max_alignment_id= cur.execute("SELECT MAX(ALIGNMENT_ID) FROM MOVIE_ALIGNMENT_LIST").fetchone()[0]
+    if max_alignment_id is None:
+        max_alignment_id = 0
+    alignment_job_id= cur.execute("SELECT MAX(ALIGNMENT_JOB_ID) FROM MOVIE_ALIGNMENT_LIST").fetchone()[0]
+    if alignment_job_id is None:
+        alignment_job_id = 1
+    else:
+        alignment_job_id += 1
+
+    max_image_asset_id= cur.execute("SELECT MAX(IMAGE_ASSET_ID) FROM IMAGE_ASSETS").fetchone()[0]
+    if max_image_asset_id is None:
+        max_image_asset_id = 0
+    for result in results:
+        movie_info = cur.execute("SELECT X_SIZE, Y_SIZE, MOVIE_ASSET_ID, NAME, PROTEIN_IS_WHITE, SPHERICAL_ABERRATION FROM MOVIE_ASSETS WHERE FILENAME = ?", (parameters[result["parameter_index"]].input_filename,)).fetchone()
+        mrc = mrcfile.open(parameters[result["parameter_index"]].output_filename)
+        xsize = mrc.header.nx
+        ysize = mrc.header.ny
+        if result["orig_x"] > 0:
+            x_bin_factor       = movie_info[0] / result["orig_x"]
+            y_bin_factor       = movie_info[1] / result["orig_y"]
+        else:
+            x_bin_factor       = movie_info[0] / xsize
+            y_bin_factor       = movie_info[1] / ysize
+        average_bin_factor = (x_bin_factor + y_bin_factor) / 2.0
+        actual_pixel_size = parameters[result["parameter_index"]].pixel_size * average_bin_factor
+        MOVIE_ALIGNMENT_LIST.append({
+                    "ALIGNMENT_ID" : max_alignment_id + 1,
+                    "DATETIME_OF_RUN" : datetime_to_msdos(datetime.datetime.now()),
+                    "ALIGNMENT_JOB_ID": alignment_job_id,
+                    "MOVIE_ASSET_ID": movie_info[2],
+                    "OUTPUT_FILE": parameters[result["parameter_index"]].output_filename,
+                    "VOLTAGE": parameters[result["parameter_index"]].acceleration_voltage,
+                    "PIXEL_SIZE": actual_pixel_size,
+                    "EXPOSURE_PER_FRAME": parameters[result["parameter_index"]].exposure_per_frame,
+                    "PRE_EXPOSURE_AMOUNT": parameters[result["parameter_index"]].pre_exposure_amount,
+                    "MIN_SHIFT": parameters[result["parameter_index"]].minimum_shift_in_angstroms,
+                    "MAX_SHIFT": parameters[result["parameter_index"]].maximum_shift_in_angstroms,
+                    "SHOULD_DOSE_FILTER": parameters[result["parameter_index"]].should_dose_filter,
+                    "SHOULD_RESTORE_POWER": parameters[result["parameter_index"]].should_restore_power,
+                    "TERMINATION_THRESHOLD": parameters[result["parameter_index"]].termination_threshold_in_angstroms,
+                    "MAX_ITERATIONS": parameters[result["parameter_index"]].max_iterations ,
+                    "BFACTOR": parameters[result["parameter_index"]].bfactor_in_angstroms,
+                    "SHOULD_MASK_CENTRAL_CROSS": parameters[result["parameter_index"]].should_mask_central_cross,
+                    "HORIZONTAL_MASK": parameters[result["parameter_index"]].horizontal_mask_size,
+                    "VERTICAL_MASK": parameters[result["parameter_index"]].vertical_mask_size,
+                    "SHOULD_INCLUDE_ALL_FRAMES_IN_SUM": True,
+                    "FIRST_FRAME_TO_SUM": parameters[result["parameter_index"]].first_frame,
+                    "LAST_FRAME_TO_SUM": parameters[result["parameter_index"]].last_frame,
+                    "ORIGINAL_X_SIZE": result["orig_x"],
+                    "ORIGINAL_Y_SIZE": result["orig_y"],
+                    "CROP_CENTER_X": result["crop_x"],
+                    "CROP_CENTER_Y": result["crop_y"],
+                    })
+        # Check if there is existing image asset
+
+        #existing_image_asset = cur.execute("SELECT IMAGE_ASSET_ID FROM IMAGE_ASSETS WHERE PARENT_MOVIE_ID = ?",(movie_info[2],)).fetchone()
+        #if existing_image_asset is None:
+        max_image_asset_id += 1
+        image_asset_id = max_image_asset_id
+        #else:
+        #    image_asset_id = existing_image_asset[0]
+
+        if change_image_assets:
+            cur.execute("REPLACE INTO IMAGE_ASSETS (IMAGE_ASSET_ID, NAME, FILENAME, POSITION_IN_STACK, PARENT_MOVIE_ID, ALIGNMENT_ID, CTF_ESTIMATION_ID, X_SIZE, Y_SIZE, PIXEL_SIZE, VOLTAGE, SPHERICAL_ABERRATION, PROTEIN_IS_WHITE, ORIGINAL_X_SIZE, ORIGINAL_Y_SIZE, CROP_CENTER_X, CROP_CENTER_Y) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)",
+                (image_asset_id, movie_info[3], parameters[result["parameter_index"]].output_filename, 1,movie_info[2],max_alignment_id+1, -1, int(xsize), int(ysize), actual_pixel_size, parameters[result["parameter_index"]].acceleration_voltage, movie_info[5], movie_info[4],result["orig_x"],result["orig_y"],result["crop_x"],result["crop_y"]))
+        MOVIE_ALIGNMENT_PARAMETERS = {
+            "FRAME_NUMBER": range(1,len(result["x_shifts"])+1),
+            "X_SHIFT": result["x_shifts"],
+            "Y_SHIFT": result["y_shifts"],
+        }
+        conn.commit()
+        MOVIE_ALIGNMENT_PARAMETERS = pd.DataFrame(MOVIE_ALIGNMENT_PARAMETERS)
+        MOVIE_ALIGNMENT_PARAMETERS.to_sql(f"MOVIE_ALIGNMENT_PARAMETERS_{max_alignment_id+1}", conn, if_exists="fail", index=False)
+        max_alignment_id += 1
+    MOVIE_ALIGNMENT_LIST = pd.DataFrame(MOVIE_ALIGNMENT_LIST)
+    MOVIE_ALIGNMENT_LIST.to_sql("MOVIE_ALIGNMENT_LIST", conn, if_exists="append", index=False)
+    conn.close()
+
+
+
+
+
+async def handle_results(reader, writer, logger):
+    #logger.info("Handling results")
+    await reader.read(4)
+    length = await reader.read(4)
+    number_of_bytes = int.from_bytes(length, byteorder="little")
+    results = await reader.read(number_of_bytes*4)
+    return(results)
+
+signal_handlers = {
+    socket_send_next_job : handle_results
+}
+
+def run(parameters: Union[UnblurParameters,list[UnblurParameters]],**kwargs):
+
+    if not isinstance(parameters, list):
+        parameters = [parameters]
+    
+    byte_results = asyncio.run(cistem_program.run("unblur", parameters, signal_handlers=signal_handlers,**kwargs))
+    result_shifts = []
+
+    for parameter_index,byte_result in byte_results:
+        number_of_images = int(((len(byte_result) /4 ) - 4 ) /2)
+        x_shifts = []
+        for offset in range(number_of_images):
+            x_shifts.append(struct.unpack_from("<f",byte_result,offset=offset*4)[0])
+        y_shifts = []
+        for offset in range(number_of_images):
+            y_shifts.append(struct.unpack_from("<f",byte_result,offset=offset*4+number_of_images*4)[0])
+        orig_x = int(struct.unpack_from("<f",byte_result,offset=2*4*number_of_images)[0])
+        orig_y = int(struct.unpack_from("<f",byte_result,offset=2*4*number_of_images+4)[0])
+        crop_x = int(struct.unpack_from("<f",byte_result,offset=2*4*number_of_images+8)[0])
+        crop_y = int(struct.unpack_from("<f",byte_result,offset=2*4*number_of_images+12)[0])
+        result_shifts.append({
+            "parameter_index": parameter_index,
+            "x_shifts": x_shifts,
+            "y_shifts": y_shifts,
+            "orig_x": orig_x,
+            "orig_y": orig_y,
+            "crop_x": crop_x,
+            "crop_y": crop_y
+        })
+
+    return(result_shifts)
+
+async def run_async(parameters: Union[UnblurParameters,list[UnblurParameters]],**kwargs):
+
+    if not isinstance(parameters, list):
+        parameters = [parameters]
+    byte_results = await cistem_program.run("unblur", parameters, signal_handlers=signal_handlers,**kwargs)
+    result_shifts = []
+
+    for parameter_index,byte_result in byte_results:
+        number_of_images = int(((len(byte_result) /4 ) - 4 ) /2)
+        x_shifts = []
+        for offset in range(number_of_images):
+            x_shifts.append(struct.unpack_from("<f",byte_result,offset=offset*4)[0])
+        y_shifts = []
+        for offset in range(number_of_images):
+            y_shifts.append(struct.unpack_from("<f",byte_result,offset=offset*4+number_of_images*4)[0])
+        orig_x = int(struct.unpack_from("<f",byte_result,offset=2*4*number_of_images)[0])
+        orig_y = int(struct.unpack_from("<f",byte_result,offset=2*4*number_of_images+4)[0])
+        crop_x = int(struct.unpack_from("<f",byte_result,offset=2*4*number_of_images+8)[0])
+        crop_y = int(struct.unpack_from("<f",byte_result,offset=2*4*number_of_images+12)[0])
+        result_shifts.append({
+            "parameter_index": parameter_index,
+            "x_shifts": x_shifts,
+            "y_shifts": y_shifts,
+            "orig_x": orig_x,
+            "orig_y": orig_y,
+            "crop_x": crop_x,
+            "crop_y": crop_y
+        })
+
+    return(result_shifts)

pycistem/programs/unblur_patch.py

@@ -0,0 +1,204 @@
+import asyncio
+import datetime
+import sqlite3
+import struct
+from dataclasses import dataclass
+from pathlib import Path
+from typing import List, Union
+
+import mrcfile
+import pandas as pd
+
+from pycistem.database import datetime_to_msdos, get_movie_info_from_db
+from pycistem.programs import cistem_program
+from pycistem.programs._cistem_constants import socket_job_result, socket_send_next_job
+
+
+@dataclass
+class UnblurPatchParameters:
+    input_filename: str
+    output_filename: str = "unblurred.mrc"
+    pixel_size: float = 1.0
+    minimum_shift_in_angstroms: float = 2.00
+    maximum_shift_in_angstroms: float = 40.0
+    should_dose_filter: bool = True
+    should_restore_power: bool = True
+    termination_threshold_in_angstroms: float = 1.0
+    max_iterations: int = 10
+    bfactor_in_angstroms: float = 1500
+    should_mask_central_cross: bool = True
+    horizontal_mask_size: int = 1
+    vertical_mask_size: int = 1
+    acceleration_voltage: float = 300.0
+    exposure_per_frame: float = 0.0
+    pre_exposure_amount: float = 0.0
+    movie_is_gain_corrected: bool = False
+    gain_filename: str = "gain.mrc"
+    movie_is_dark_corrected: bool = True
+    dark_filename: str = "dark.mrc"
+    output_binning_factor: float = 2.0
+    correct_mag_distortion: bool = False
+    mag_distortion_angle: float = 0.0
+    mag_distortion_major_scale: float = 1.0
+    mag_distortion_minor_scale: float = 1.0
+    write_out_amplitude_spectrum: bool = True
+    amplitude_spectrum_filename: str = "amplitude_spectrum.mrc"
+    write_out_small_sum_image: bool = True
+    small_sum_image_filename: str = "scaled_sum.mrc"
+    first_frame: int = 1
+    last_frame: int = 0
+    number_of_frames_for_running_average: int = 1
+    max_threads: int = 1
+    save_aligned_frames: bool = False
+    eer_frames_per_image: int = 0
+    eer_super_res_factor: int = 1
+    outputpath: str = "/tmp"
+    patchcorrection: bool = False
+    override_patchnum: bool = False
+    patch_num_x: int = 6
+    patch_num_y: int = 4
+    distortion_model: int = 3
+
+def parameters_from_database(database, decolace=False, **kwargs):
+    movie_info = get_movie_info_from_db(database)
+    ProjectDirectory = Path(database).parent
+    par = [UnblurPatchParameters(
+        input_filename = movie["FILENAME"],
+        output_filename= (ProjectDirectory / "Assets" / "Images" / f"{Path(movie['FILENAME']).stem}_{movie['MOVIE_ASSET_ID']}_auto.mrc").as_posix(),
+        pixel_size = movie["PIXEL_SIZE"],
+        gain_filename=movie["GAIN_FILENAME"],
+        output_binning_factor=movie["OUTPUT_BINNING_FACTOR"],
+        exposure_per_frame=movie["DOSE_PER_FRAME"],
+        amplitude_spectrum_filename=(ProjectDirectory / "Assets" / "Images" / "Spectra" / f"{Path(movie['FILENAME']).stem}_{movie['MOVIE_ASSET_ID']}_auto.mrc").as_posix(),
+        small_sum_image_filename=(ProjectDirectory / "Assets" / "Images" / "Scaled" / f"{Path(movie['FILENAME']).stem}_{movie['MOVIE_ASSET_ID']}_auto.mrc").as_posix(),
+    ) for i,movie in movie_info.iterrows()]
+    return(par)
+
+def write_results_to_database(database,  parameters, results,change_image_assets):
+    conn = sqlite3.connect(database)
+    cur = conn.cursor()
+    results = sorted(results, key=lambda x: x["parameter_index"])
+    MOVIE_ALIGNMENT_LIST = []
+
+    max_alignment_id= cur.execute("SELECT MAX(ALIGNMENT_ID) FROM MOVIE_ALIGNMENT_LIST").fetchone()[0]
+    if max_alignment_id is None:
+        max_alignment_id = 0
+    alignment_job_id= cur.execute("SELECT MAX(ALIGNMENT_JOB_ID) FROM MOVIE_ALIGNMENT_LIST").fetchone()[0]
+    if alignment_job_id is None:
+        alignment_job_id = 1
+    else:
+        alignment_job_id += 1
+
+    max_image_asset_id= cur.execute("SELECT MAX(IMAGE_ASSET_ID) FROM IMAGE_ASSETS").fetchone()[0]
+    if max_image_asset_id is None:
+        max_image_asset_id = 0
+    for result in results:
+        movie_info = cur.execute("SELECT X_SIZE, Y_SIZE, MOVIE_ASSET_ID, NAME, PROTEIN_IS_WHITE, SPHERICAL_ABERRATION FROM MOVIE_ASSETS WHERE FILENAME = ?", (parameters[result["parameter_index"]].input_filename,)).fetchone()
+        x_bin_factor       = movie_info[0] / result["orig_x"]
+        y_bin_factor       = movie_info[1] / result["orig_y"]
+        average_bin_factor = (x_bin_factor + y_bin_factor) / 2.0
+        actual_pixel_size = parameters[result["parameter_index"]].pixel_size * average_bin_factor
+        MOVIE_ALIGNMENT_LIST.append({
+                    "ALIGNMENT_ID" : max_alignment_id + 1,
+                    "DATETIME_OF_RUN" : datetime_to_msdos(datetime.datetime.now()),
+                    "ALIGNMENT_JOB_ID": alignment_job_id,
+                    "MOVIE_ASSET_ID": movie_info[2],
+                    "OUTPUT_FILE": parameters[result["parameter_index"]].output_filename,
+                    "VOLTAGE": parameters[result["parameter_index"]].acceleration_voltage,
+                    "PIXEL_SIZE": actual_pixel_size,
+                    "EXPOSURE_PER_FRAME": parameters[result["parameter_index"]].exposure_per_frame,
+                    "PRE_EXPOSURE_AMOUNT": parameters[result["parameter_index"]].pre_exposure_amount,
+                    "MIN_SHIFT": parameters[result["parameter_index"]].minimum_shift_in_angstroms,
+                    "MAX_SHIFT": parameters[result["parameter_index"]].maximum_shift_in_angstroms,
+                    "SHOULD_DOSE_FILTER": parameters[result["parameter_index"]].should_dose_filter,
+                    "SHOULD_RESTORE_POWER": parameters[result["parameter_index"]].should_restore_power,
+                    "TERMINATION_THRESHOLD": parameters[result["parameter_index"]].termination_threshold_in_angstroms,
+                    "MAX_ITERATIONS": parameters[result["parameter_index"]].max_iterations ,
+                    "BFACTOR": parameters[result["parameter_index"]].bfactor_in_angstroms,
+                    "SHOULD_MASK_CENTRAL_CROSS": parameters[result["parameter_index"]].should_mask_central_cross,
+                    "HORIZONTAL_MASK": parameters[result["parameter_index"]].horizontal_mask_size,
+                    "VERTICAL_MASK": parameters[result["parameter_index"]].vertical_mask_size,
+                    "SHOULD_INCLUDE_ALL_FRAMES_IN_SUM": True,
+                    "FIRST_FRAME_TO_SUM": parameters[result["parameter_index"]].first_frame,
+                    "LAST_FRAME_TO_SUM": parameters[result["parameter_index"]].last_frame,
+                    "ORIGINAL_X_SIZE": result["orig_x"],
+                    "ORIGINAL_Y_SIZE": result["orig_y"],
+                    "CROP_CENTER_X": result["crop_x"],
+                    "CROP_CENTER_Y": result["crop_y"],
+                    })
+        # Check if there is existing image asset
+
+        existing_image_asset = cur.execute("SELECT IMAGE_ASSET_ID FROM IMAGE_ASSETS WHERE PARENT_MOVIE_ID = ?",(movie_info[2],)).fetchone()
+        if existing_image_asset is None:
+            max_image_asset_id += 1
+            image_asset_id = max_image_asset_id
+        else:
+            image_asset_id = existing_image_asset[0]
+
+        mrc = mrcfile.open(parameters[result["parameter_index"]].output_filename)
+        xsize = mrc.header.nx
+        ysize = mrc.header.ny
+        cur.execute("REPLACE INTO IMAGE_ASSETS (IMAGE_ASSET_ID, NAME, FILENAME, POSITION_IN_STACK, PARENT_MOVIE_ID, ALIGNMENT_ID, CTF_ESTIMATION_ID, X_SIZE, Y_SIZE, PIXEL_SIZE, VOLTAGE, SPHERICAL_ABERRATION, PROTEIN_IS_WHITE, ORIGINAL_X_SIZE, ORIGINAL_Y_SIZE, CROP_CENTER_X, CROP_CENTER_Y) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)",
+            (image_asset_id, movie_info[3], parameters[result["parameter_index"]].output_filename, 1,movie_info[2],max_alignment_id+1, -1, int(xsize), int(ysize), actual_pixel_size, parameters[result["parameter_index"]].acceleration_voltage, movie_info[5], movie_info[4],result["orig_x"],result["orig_y"],result["crop_x"],result["crop_y"]))
+        MOVIE_ALIGNMENT_PARAMETERS = {
+            "FRAME_NUMBER": range(1,len(result["x_shifts"])+1),
+            "X_SHIFT": result["x_shifts"],
+            "Y_SHIFT": result["y_shifts"],
+        }
+        conn.commit()
+        MOVIE_ALIGNMENT_PARAMETERS = pd.DataFrame(MOVIE_ALIGNMENT_PARAMETERS)
+        MOVIE_ALIGNMENT_PARAMETERS.to_sql(f"MOVIE_ALIGNMENT_PARAMETERS_{max_alignment_id+1}", conn, if_exists="fail", index=False)
+        max_alignment_id += 1
+    MOVIE_ALIGNMENT_LIST = pd.DataFrame(MOVIE_ALIGNMENT_LIST)
+    MOVIE_ALIGNMENT_LIST.to_sql("MOVIE_ALIGNMENT_LIST", conn, if_exists="append", index=False)
+    conn.close()
+
+
+
+
+
+async def handle_results(reader, writer, logger):
+    #logger.info("Handling results")
+    await reader.read(4)
+    length = await reader.read(4)
+    number_of_bytes = int.from_bytes(length, byteorder="little")
+    results = await reader.read(number_of_bytes*4)
+    print(struct.unpack_from("<"+''.join(["f" for a in range(number_of_bytes)]),results))
+    return(results)
+
+signal_handlers = {
+    socket_send_next_job : handle_results
+}
+
+def run(parameters: Union[UnblurPatchParameters,list[UnblurPatchParameters]], unblur_command: str="unblur", **kwargs):
+
+    if not isinstance(parameters, list):
+        parameters = [parameters]
+
+    byte_results = asyncio.run(cistem_program.run(unblur_command, parameters, signal_handlers=signal_handlers,**kwargs))
+    result_shifts = []
+
+    for parameter_index,byte_result in byte_results:
+        number_of_images = int(((len(byte_result) /4 ) - 4 ) /2)
+        x_shifts = []
+        for offset in range(number_of_images):
+            x_shifts.append(struct.unpack_from("<f",byte_result,offset=offset*4)[0])
+        y_shifts = []
+        for offset in range(number_of_images):
+            y_shifts.append(struct.unpack_from("<f",byte_result,offset=offset*4+number_of_images*4)[0])
+        orig_x = int(struct.unpack_from("<f",byte_result,offset=2*4*number_of_images)[0])
+        orig_y = int(struct.unpack_from("<f",byte_result,offset=2*4*number_of_images+4)[0])
+        crop_x = int(struct.unpack_from("<f",byte_result,offset=2*4*number_of_images+8)[0])
+        crop_y = int(struct.unpack_from("<f",byte_result,offset=2*4*number_of_images+12)[0])
+        result_shifts.append({
+            "parameter_index": parameter_index,
+            "x_shifts": x_shifts,
+            "y_shifts": y_shifts,
+            "orig_x": orig_x,
+            "orig_y": orig_y,
+            "crop_x": crop_x,
+            "crop_y": crop_y
+        })
+
+
+    return(result_shifts)

pycistem/utils/__init__.py

@@ -0,0 +1 @@
+from .extract_particles import extract_particles

pycistem/utils/extract_particles.py

@@ -0,0 +1,36 @@
+from pathlib import Path
+
+def extract_particles(starfile_filename: Path, stack_filename: Path, box_size: int = 256):
+    """Extract particles from a star file and save them as individual images.
+
+    Args:
+        starfile_filename (str): The filename of the star file.
+        box_size (int, optional): The size of the extracted particles. Defaults to 256.
+    """
+    import starfile
+    import mrcfile
+    import numpy as np
+    from itertools import groupby
+    particle_info = starfile.read(starfile_filename)
+    mrc = mrcfile.new_mmap(stack_filename, (len(particle_info), box_size, box_size), mrc_mode=2, overwrite=True)
+    # Iterate over groupby cisTEMOriginalImageFilename
+
+    for micrograph_filename, subparticles in groupby(particle_info.itertuples(), lambda x: x.cisTEMOriginalImageFilename):
+        micrograph = mrcfile.open(micrograph_filename)
+        if micrograph.data.ndim == 3:
+            micrograph_data  = micrograph.data[0].copy()
+        else:
+            micrograph_data = micrograph.data.copy()
+        for particle in subparticles:
+            x = round(particle.cisTEMOriginalXPosition/particle.cisTEMPixelSize)
+            y = round(particle.cisTEMOriginalYPosition/particle.cisTEMPixelSize)
+            particle_image = micrograph_data.take(range(y-box_size//2, y+box_size//2), mode='clip', axis=0).take(range(x-box_size//2, x+box_size//2), mode='clip', axis=1)
+            if particle_image.shape != (box_size, box_size):
+                raise ValueError(f"Particle at {x},{y} from micrograph {micrograph_filename} {micrograph_data.shape} is out of bounds {particle_image.shape}.")
+            
+            particle_image -= particle_image.mean()
+            particle_image /= particle_image.std()
+            mrc.data[particle.cisTEMPositionInStack-1] = particle_image
+            yield
+    mrc.close()
+    return

pycistem/utils/move_class_into_unbinned_particle_stack.py

@@ -0,0 +1,34 @@
+import pandas as pd
+from typing import Union
+import sqlite3
+import typer
+from pathlib import Path
+from typing_extensions import Annotated
+import starfile
+
+app = typer.Typer()
+
+
+@app.command()
+def plot(database: Annotated[Path, typer.Argument(...,help="The database file to use")],
+        refinement_id: Annotated[int, typer.Argument(...,help="The refinement to use")],
+        class_id: Annotated[int, typer.Argument(...,help="The class to cus")],
+        input_star_file: Annotated[Path, typer.Argument(...,help="The starfile to use")],
+        output_star_file: Annotated[Path, typer.Argument(...,help="The starfile to write to")],
+        ):
+    db = sqlite3.connect(database)
+    refinements = []
+    refinment_info = pd.read_sql_query(f"SELECT * FROM REFINEMENT_LIST WHERE REFINEMENT_ID = {refinement_id}", db).iloc[0]
+    refinment_result = pd.read_sql_query(f"SELECT * FROM REFINEMENT_RESULT_{refinement_id}_{class_id}", db)
+    refinment_package_info = pd.read_sql_query(f"SELECT * FROM REFINEMENT_PACKAGE_CONTAINED_PARTICLES_{refinment_info['REFINEMENT_PACKAGE_ASSET_ID']}", db)
+    
+    starfile_info = starfile.read(input_star_file)
+
+    original_ids = refinment_package_info["ORIGINAL_PARTICLE_POSITION_ASSET_ID"].to_list()
+    subset_starfile_info = starfile_info.iloc[original_ids].copy()
+    subset_starfile_info["cisTEMOccupancy"] = refinment_result["OCCUPANCY"].to_list()
+    subset_starfile_info["cisTEMScore"] = refinment_result["SCORE"].to_list()
+    starfile.write(subset_starfile_info, output_star_file)
+
+if __name__ == "__main__":
+    app()

pycistem/utils/order_by_class_occupancy.py

@@ -0,0 +1,26 @@
+import pandas as pd
+from typing import Union
+import sqlite3
+import typer
+from pathlib import Path
+from typing_extensions import Annotated
+import matplotlib.pyplot as plt
+
+app = typer.Typer()
+
+
+@app.command()
+def plot(database: Annotated[Path, typer.Argument(...,help="The database file to use")],
+        start_refinement: Annotated[int, typer.Argument(...,help="The refinement to start from")]):
+    db = sqlite3.connect(database)
+    refinment_info = pd.read_sql_query(f"SELECT * FROM REFINEMENT_LIST WHERE REFINEMENT_ID = {start_refinement}", db)
+    class_info = pd.read_sql_query(f"SELECT * FROM REFINEMENT_DETAILS_{refinment_info['REFINEMENT_ID'].values[0]}", db)
+    class_info.sort_values(by="AVERAGE_OCCUPANCY", inplace=True)
+    # print as a table the CLASS_NUMBER and AVERAGE_OCCUPANCY
+    print(class_info[['CLASS_NUMBER', 'AVERAGE_OCCUPANCY']])
+
+        
+            
+
+if __name__ == "__main__":
+    app()

pycistem/utils/plot_class_occupancy.py

@@ -0,0 +1,51 @@
+import pandas as pd
+from typing import Union
+import sqlite3
+import typer
+from pathlib import Path
+from typing_extensions import Annotated
+import matplotlib.pyplot as plt
+
+app = typer.Typer()
+
+def return_class_occupancies(row, db):
+    class_info = pd.read_sql_query(f"SELECT * FROM REFINEMENT_DETAILS_{row['REFINEMENT_ID']}", db)
+    return [row['REFINEMENT_ID']] + class_info['AVERAGE_OCCUPANCY'].to_list()
+
+def return_num_part_with_occ_higher_than(row, db, min_occ):
+    class_info = pd.read_sql_query(f"SELECT * FROM REFINEMENT_DETAILS_{row['REFINEMENT_ID']}", db)
+    result = [row['REFINEMENT_ID']]
+    for class_id in class_info['CLASS_NUMBER']:
+        particle_info = pd.read_sql_query(f"SELECT OCCUPANCY FROM REFINEMENT_RESULT_{row['REFINEMENT_ID']}_{class_id}", db)
+        result.append(len(particle_info[particle_info['OCCUPANCY'] > min_occ]))
+    print(result)
+    return result
+
+@app.command()
+def plot(database: Annotated[Path, typer.Argument(...,help="The database file to use")],
+        start_refinement: Annotated[int, typer.Argument(...,help="The refinement to start from")]):
+    db = sqlite3.connect(database)
+    refinements = []
+    refinment_info = pd.read_sql_query(f"SELECT * FROM REFINEMENT_LIST WHERE REFINEMENT_ID = {start_refinement}", db)
+    while len(refinment_info) == 1:
+        reference_refinement_id = refinment_info["REFINEMENT_ID"].values[0]
+        refinements.append(refinment_info)
+        refinment_info = pd.read_sql_query(f"SELECT * FROM REFINEMENT_LIST WHERE STARTING_REFINEMENT_ID = {reference_refinement_id}", db)
+    refinements = pd.concat(refinements)
+    av_occs = refinements.apply(return_class_occupancies, axis=1, result_type='expand', args=(db,))
+    av_occs.plot(x=0)
+    plt.savefig("average_occupancies.png")
+    plt.figure()
+    occ_0_5 = refinements.apply(return_num_part_with_occ_higher_than, axis=1, result_type='expand', args=(db, 50))
+    occ_0_5.plot(x=0)
+    plt.savefig("occupancies_0_5.png")
+    plt.figure()
+    occ_0_9 = refinements.apply(return_num_part_with_occ_higher_than, axis=1, result_type='expand', args=(db, 90))
+    occ_0_9.plot(x=0)
+    plt.savefig("occupancies_0_9.png")
+
+        
+            
+
+if __name__ == "__main__":
+    app()