From 019c66bb87b71f3761448cb6336ad0a63853400c Mon Sep 17 00:00:00 2001 From: kekexin714 Date: Thu, 11 Sep 2025 10:27:26 -0400 Subject: [PATCH] Add python package for 2DTM postprocessing --- 2DTM_postprocess_tool/README.md | 59 +++++++ 2DTM_postprocess_tool/setup.py | 25 +++ 2DTM_postprocess_tool/src/cli/__init__.py | 0 .../src/cli/compare_starfiles.py | 42 +++++ .../src/cli/extract_particles.py | 65 +++++++ .../src/cli/filter_particles.py | 121 +++++++++++++ 2DTM_postprocess_tool/src/cli/update_par.py | 70 ++++++++ 2DTM_postprocess_tool/src/tm_post/__init__.py | 0 .../src/tm_post/compare_starfiles.py | 70 ++++++++ 2DTM_postprocess_tool/src/tm_post/database.py | 148 ++++++++++++++++ 2DTM_postprocess_tool/src/tm_post/extract.py | 136 ++++++++++++++ 2DTM_postprocess_tool/src/tm_post/filters.py | 131 ++++++++++++++ 2DTM_postprocess_tool/src/tm_post/geodesic.py | 92 ++++++++++ 2DTM_postprocess_tool/src/tm_post/geometry.py | 60 +++++++ .../src/tm_post/image_data.py | 25 +++ 2DTM_postprocess_tool/src/tm_post/mrcfile.py | 8 + 2DTM_postprocess_tool/src/tm_post/peak.py | 50 ++++++ 2DTM_postprocess_tool/src/tm_post/starfile.py | 166 ++++++++++++++++++ .../src/tm_post/statistics.py | 117 ++++++++++++ 19 files changed, 1385 insertions(+) create mode 100644 2DTM_postprocess_tool/README.md create mode 100644 2DTM_postprocess_tool/setup.py create mode 100644 2DTM_postprocess_tool/src/cli/__init__.py create mode 100644 2DTM_postprocess_tool/src/cli/compare_starfiles.py create mode 100644 2DTM_postprocess_tool/src/cli/extract_particles.py create mode 100644 2DTM_postprocess_tool/src/cli/filter_particles.py create mode 100644 2DTM_postprocess_tool/src/cli/update_par.py create mode 100644 2DTM_postprocess_tool/src/tm_post/__init__.py create mode 100644 2DTM_postprocess_tool/src/tm_post/compare_starfiles.py create mode 100644 2DTM_postprocess_tool/src/tm_post/database.py create mode 100644 2DTM_postprocess_tool/src/tm_post/extract.py create mode 100644 2DTM_postprocess_tool/src/tm_post/filters.py create mode 100644 2DTM_postprocess_tool/src/tm_post/geodesic.py create mode 100644 2DTM_postprocess_tool/src/tm_post/geometry.py create mode 100644 2DTM_postprocess_tool/src/tm_post/image_data.py create mode 100644 2DTM_postprocess_tool/src/tm_post/mrcfile.py create mode 100644 2DTM_postprocess_tool/src/tm_post/peak.py create mode 100644 2DTM_postprocess_tool/src/tm_post/starfile.py create mode 100644 2DTM_postprocess_tool/src/tm_post/statistics.py diff --git a/2DTM_postprocess_tool/README.md b/2DTM_postprocess_tool/README.md new file mode 100644 index 000000000..6a74be183 --- /dev/null +++ b/2DTM_postprocess_tool/README.md @@ -0,0 +1,59 @@ +# 2DTM Postprocessing + +A modular Python package for postprocessing 2D template matching results from cryo-EM workflows (e.g., cisTEM), including 2DTM p-value calculation, particle extraction and filtering. + +--- + +## Installation + +```bash +git clone https://github.com/kekexinz/2DTM_postprocess_tool.git +cd 2DTM_postprocess_tool +pip install -e . # editable mode +``` + +## 📦 Usage + +### `extract-particles` +Extract initial particle peaks from 2DTM search. +```bash +extract-particles \ +--db_file \ +--tm_job_id 1 \ +--ctf_job_id 1 \ +--pixel_size 1.0 \ +--output +[--metric pval] \ # "zscore" or "pval" +[--metric_cutoff 8.0] \ +[--threads 22] \ +[--local_max_filter] \ # "snr" or "zscore" (default) used for skimage peak_local_max +[--min_peak_radius 10] \ # used for "min_distance" in skimage peak_local_max +[--exclude_borders 92] \ # avoid finding partial particles near the edge of the image, used for skimage peak_local_max +[--quadrants 1] \ # 1 (default) or 3, calculating p-value for only the first-quadrant or quadrant 1,2,4 (recommended for small particles) + +``` + +### `filter-particles` + +Filter particles based on image thickness and/or angular invariance. + +```bash +filter-particles \ + --star_file \ # output from extract-particles + --db_file \ + --tm_job_id 1 \ + --ctf_job_id 1 \ + --pixel_size 1.0 \ + --output filtered_peaks.star \ + [--avg_cutoff_lb] \ # angular search CC per-pixel avg + [--sd_cutoff_ub] \ # angular search CC per-pixel sd + [--snr_cutoff_ub] \ + [--filter_by_image_thickness] \ # ctffind5 parameters + [--thickness_cutoff_lb] \ + [--thickness_cutoff_ub] \ + [--ctf_fitting_score_lb] \ + [--ctf_fitting_score_ub] \ +``` + +### 3D reconstruction & refinement in cisTEM +The output extracted_peaks.star and filtered_peaks.star can be imported into cisTEM as a RefinementPackage for further 3D reconstruction and refinement. diff --git a/2DTM_postprocess_tool/setup.py b/2DTM_postprocess_tool/setup.py new file mode 100644 index 000000000..d2a7ac016 --- /dev/null +++ b/2DTM_postprocess_tool/setup.py @@ -0,0 +1,25 @@ +from setuptools import setup, find_packages + +setup( + name="tm_post", + version="0.1", + author="Kexin Zhang", + description="Postprocessing utilities for 2D template matching", + packages=find_packages(where="src"), + package_dir={"": "src"}, + install_requires=[ + "numpy", + "pandas", + "scipy", + "mrcfile", + "joblib", + "tqdm", + "scikit-image", + ], + entry_points={ + "console_scripts": [ + "filter-particles = cli.filter_particles:main", + "extract-particles = cli.extract_particles:main", + ], + }, +) \ No newline at end of file diff --git a/2DTM_postprocess_tool/src/cli/__init__.py b/2DTM_postprocess_tool/src/cli/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/2DTM_postprocess_tool/src/cli/compare_starfiles.py b/2DTM_postprocess_tool/src/cli/compare_starfiles.py new file mode 100644 index 000000000..179a65279 --- /dev/null +++ b/2DTM_postprocess_tool/src/cli/compare_starfiles.py @@ -0,0 +1,42 @@ +import argparse +import pandas as pd +import tm_post.starfile as starfile +from tm_post.compare_starfiles import compare_starfiles_for_matched_peaks + +def parse_arguments(): + parser = argparse.ArgumentParser( + description="Compare two STAR files and extract matched peaks based on spatial and angular thresholds." + ) + parser.add_argument('--starfile_a', type=str, required=True, help="Path to first STAR file (e.g. bin2x).") + parser.add_argument('--starfile_b', type=str, required=True, help="Path to second STAR file (e.g. bin1x).") + parser.add_argument('--d_xy_cutoff', type=float, default=10.0, help="Maximum XY distance in Å for matching.") + parser.add_argument('--euler_err_cutoff', type=float, default=5.0, help="Maximum Euler angle error in degrees.") + parser.add_argument('--pattern', type=str, default=r"mc2_[12]x_(.*?frames)", help="Regex pattern for extracting match key.") + parser.add_argument('--output', type=str, required=True, help="Path to output STAR file with matched peaks.") + + return parser.parse_args() + +def main(): + args = parse_arguments() + + print("[INFO] Comparing starfiles...") + matched_df = compare_starfiles_for_matched_peaks( + starfile_a=args.starfile_a, + starfile_b=args.starfile_b, + d_xy_cutoff=args.d_xy_cutoff, + euler_err_cutoff=args.euler_err_cutoff, + pattern=args.pattern + ) + + if matched_df.empty: + print("[INFO] No matching peaks found.") + return + + # Convert matched_df to STAR format (with dummy column) and write with standard header + matched_df_star = starfile.add_star_dummy_column(matched_df) + header_lines = starfile.read_tm_package_starfile_header() + starfile.write_starfile_with_headers(args.output, header_lines, matched_df_star) + print(f"[INFO] Matched particles saved to: {args.output}") + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/2DTM_postprocess_tool/src/cli/extract_particles.py b/2DTM_postprocess_tool/src/cli/extract_particles.py new file mode 100644 index 000000000..ee2b9d767 --- /dev/null +++ b/2DTM_postprocess_tool/src/cli/extract_particles.py @@ -0,0 +1,65 @@ +import argparse +import pandas as pd +import tm_post.database as db +from tm_post.database import load_tm_images_from_db +from tm_post.extract import extract_particles_from_2dtm_search +from tm_post.starfile import write_starfile_with_headers, read_tm_package_starfile_header + +def parse_arguments(): + parser = argparse.ArgumentParser(description="Extract peaks from 2DTM searches.") + + parser.add_argument('--db_file', type=str, required=True, help="Path to the database file.") + parser.add_argument('--tm_job_id', type=int, required=True, help="Template match job ID.") + parser.add_argument('--ctf_job_id', type=int, required=True, help="CTF job ID.") + + parser.add_argument('--min_peak_radius', type=int, default=10, help="Cutoff for XY distance.") + parser.add_argument('--exclude_borders', type=int, default=35, help="Exclude borders in the image.") + + parser.add_argument('--local_max_filter', type=str, default="zscore", choices=["zscore", "snr"], help="Local max filter to use.") + parser.add_argument('--metric', type=str, default="pval", choices=["pval", "zscore", "snr"], help="Metric to use for filtering.") + parser.add_argument('--metric_cutoff', type=float, default=8.0, help="Selected metric cutoff.") + parser.add_argument('--pixel_size', type=float, required=True, default=1.0, help="Wanted pixel size in final stack.") + parser.add_argument('--threads', type=int, default=4, help="Number of threads for parallel processing.") + + parser.add_argument('--quadrants', type=int, default=1, help="Number of quadrants to use for filtering.") + + parser.add_argument('--output', type=str, required=True, help="Path to the output star file.") + + return parser.parse_args() + +def main(): + args = parse_arguments() + + # Load database information + print("[INFO] Loading TM image data from database...") + tm_images, df_ctf, df_info = load_tm_images_from_db( + db_file=args.db_file, + tm_job_id=args.tm_job_id, + ctf_job_id=args.ctf_job_id + ) + + print(f"[INFO] Running extraction on {len(tm_images)} images...") + + df_star = extract_particles_from_2dtm_search( + tm_images=tm_images, + local_max_filter=args.local_max_filter, + metric=args.metric, + metric_cutoff=args.metric_cutoff, + pixel_size=args.pixel_size, + max_threads=args.threads, + df_ctf=df_ctf, + df_info=df_info, + ctf_job_id=args.ctf_job_id, + min_radius=args.min_peak_radius, + exclude_borders=args.exclude_borders, + q=args.quadrants + ) + + print("[INFO] Writing STAR file...") + header_lines = read_tm_package_starfile_header() # provide default STAR header + write_starfile_with_headers(args.output, header_lines, df_star) + print(f"[INFO] Done. Extracted particles saved to {args.output}") + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/2DTM_postprocess_tool/src/cli/filter_particles.py b/2DTM_postprocess_tool/src/cli/filter_particles.py new file mode 100644 index 000000000..4528e65ce --- /dev/null +++ b/2DTM_postprocess_tool/src/cli/filter_particles.py @@ -0,0 +1,121 @@ +import argparse +import tm_post.starfile as starfile +import tm_post.database as db +from tm_post.filters import apply_filter + +def parse_arguments(): + parser = argparse.ArgumentParser(description="Filter peaks using TM results and image quality.") + # read job information from database file + parser.add_argument('--star_file', type=str, required=True, help="Path to the particle starfile.") + parser.add_argument('--db_file', type=str, required=True, help="Path to the database file.") + parser.add_argument('--tm_job_id', type=int, required=True, help="Template match job ID.") + parser.add_argument('--ctf_job_id', type=int, required=True, help="CTF job ID.") + parser.add_argument('--pixel_size', type=float, required=True, help="Pixel size in Angstroms.") + + # read particle information from .star file (extract_peaks.py output) + parser.add_argument('--avg_cutoff_lb', type=float, default=None, help="Lower bound for average cutoff.") + parser.add_argument('--sd_cutoff_ub', type=float, default=None, help="Upper bound for SD cutoff.") + parser.add_argument('--pval_cutoff_lb', type=float, default=None, help="Lower bound for p-value cutoff.") + parser.add_argument('--snr_cutoff_ub', type=float, default=None, help="Upper bound for SNR (optional).") + parser.add_argument('--snr_cutoff_lb', type=float, default=None, help="Lower bound for SNR (optional).") + parser.add_argument('--filter_by_image_thickness', action="store_true", help="Use thickness to filter good micrographs? (default: False)") + parser.add_argument('--thickness_cutoff_lb', type=float, default=None, help="Lower bound for thickness cutoff (A).") + parser.add_argument('--thickness_cutoff_ub', type=float, default=None, help="Upper bound for thickness cutoff (A).") + parser.add_argument('--filter_by_angular_invariance', action="store_true", help="Use angular invariance to filter good particles (default: False)?") + parser.add_argument('--geodesic_r', type=int, default=None, help="Radius in pixels for local patch.") + parser.add_argument('--geodesic_threads', type=int, default=None, help="Number of threads for geodesic computation.") + parser.add_argument('--geodesic_method', type=str, default=None, help="Method for geodesic filtering ('quantile' or 'cutoff').") + parser.add_argument('--geodesic_threshold', type=float, default=None, help="Threshold value for geodesic filtering.") + + parser.add_argument('--ctf_fitting_score_lb', type=float, default=None, help="Lower bound for CTF fitting score.") + parser.add_argument('--ctf_fitting_score_ub', type=float, default=None, help="Upper bound for CTF fitting score.") + + + parser.add_argument('--output', type=str, required=True, help="Path to the output star file.") + return parser.parse_args() + + +def main(): + args = parse_arguments() + + # Load particle information from star file + print("[INFO] Loading peak file...") + df_peaks = starfile.load_particle_starfile(args.star_file) + + # Extract header lines + #header_lines, _ = starfile.extract_header_lines(args.star_file) + + # Load database information + print("[INFO] Loading database...") + result = db.get_info_from_cistem_database( + args.db_file, args.tm_job_id, args.ctf_job_id + ) + + # Extract relevant data + image_list = result["image_list"] + psi_list = result["PSI_OUTPUT_FILE"] + theta_list = result["THETA_OUTPUT_FILE"] + phi_list = result["PHI_OUTPUT_FILE"] + df_ctf = result["df_ctf"] + df_info = result["df_info"] + + # Apply filters + if args.filter_by_image_thickness: + if args.thickness_cutoff_lb is None: + args.thickness_cutoff_lb = 0.0 + if args.thickness_cutoff_ub is None: + args.thickness_cutoff_ub = 500.0 + + if args.filter_by_angular_invariance: + if args.geodesic_r is None: + args.geodesic_r = 4 + if args.geodesic_threads is None: + args.geodesic_threads = 8 + if args.geodesic_method is None: + args.geodesic_method = "quantile" + if args.geodesic_threshold is None: + args.geodesic_threshold = 0.8 + + filtered_df, all_df = apply_filter( + df=df_peaks, + image_list=image_list, + psi_list=psi_list, + theta_list=theta_list, + phi_list=phi_list, + pixel_size=args.pixel_size, + df_ctf=df_ctf, + df_info=df_info, + avg_cutoff_lb=args.avg_cutoff_lb, + sd_cutoff_ub=args.sd_cutoff_ub, + pval_cutoff_lb=args.pval_cutoff_lb, + snr_cutoff_ub=args.snr_cutoff_ub, + snr_cutoff_lb=args.snr_cutoff_lb, + filter_by_image_thickness=args.filter_by_image_thickness, + thickness_lb=args.thickness_cutoff_lb, + thickness_ub=args.thickness_cutoff_ub, + ctf_fitting_score_lb=args.ctf_fitting_score_lb, + ctf_fitting_score_ub=args.ctf_fitting_score_ub, + filter_by_angular_invariance=args.filter_by_angular_invariance, + geodesic_r=args.geodesic_r, + geodesic_threads=args.geodesic_threads, + geodesic_method=args.geodesic_method, + geodesic_threshold=args.geodesic_threshold + ) + + # Save filtered STAR file with updated SCORE (no extra metadata columns) + columns_to_keep = ["ORIGINAL_IMAGE_FILENAME","ORIGX","ORIGY","AVG","SD","PVALUE","ZSCORE","SNR"] + meta_df = all_df[columns_to_keep].copy() + + # Convert filtered DataFrame to lines with empty column + starfile_df_star = starfile.add_star_dummy_column(filtered_df) + + # Write output with headers + header_lines = starfile.read_tm_package_starfile_header() # provide default STAR header + starfile.write_starfile_with_headers(args.output, header_lines, starfile_df_star) + print(f"[INFO] Filtered data saved to {args.output}") + + metadata_file = args.output.replace(".star", "_metadata.txt") + meta_df.to_csv(metadata_file, sep="\t", index=False, float_format="%.2f") + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/2DTM_postprocess_tool/src/cli/update_par.py b/2DTM_postprocess_tool/src/cli/update_par.py new file mode 100644 index 000000000..f8338ff3c --- /dev/null +++ b/2DTM_postprocess_tool/src/cli/update_par.py @@ -0,0 +1,70 @@ +#!/usr/bin/env python3 + +import argparse +import pandas as pd +from io import StringIO +from tm_post import starfile + +def read_par_file(par_path): + """Read a cisTEM .par file into header, data DataFrame, and footer.""" + with open(par_path, "r") as f: + lines = f.readlines() + + header = lines[0] + footer = lines[-2:] + data_lines = lines[1:-2] + + data_str = ''.join(data_lines) + df = pd.read_csv(StringIO(data_str), delim_whitespace=True, header=None) + df.columns = header.strip().split() + + return header, df, footer + + +def read_score_file(score_path): + """Read a file with one score per line.""" + df = starfile.load_particle_starfile(score_path) + return df["SCORE"].tolist() + + +def update_scores(df, score_files): + """Concatenate scores from all files and assign to the SCORE column.""" + all_scores = [] + for path in score_files: + scores = read_score_file(path) + all_scores.extend(scores) + + if len(all_scores) != len(df): + raise ValueError(f"Number of scores ({len(all_scores)}) does not match number of particles ({len(df)}).") + + df['SCORE'] = all_scores + return df + + +def write_par_file(out_path, header, df, footer): + """Write the updated .par file.""" + with open(out_path, "w") as f: + f.write(header) + for row in df.itertuples(index=False): + values = ' '.join(f"{v:>8}" if isinstance(v, float) else f"{v:>8}" for v in row) + f.write(f"{values}\n") + f.writelines(footer) + + +def main(): + parser = argparse.ArgumentParser(description="Update SCORE column in a cisTEM .par file.") + parser.add_argument("par_file", help="Path to the original .par file") + parser.add_argument("score_files", nargs='+', help="One or more star files containing updated SCORE values") + parser.add_argument("-o", "--output", required=True, help="Output path for updated .par file") + + args = parser.parse_args() + + header, df, footer = read_par_file(args.par_file) + df = update_scores(df, args.score_files) + write_par_file(args.output, header, df, footer) + + print(f"[INFO] Updated .par file written to: {args.output}") + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/2DTM_postprocess_tool/src/tm_post/__init__.py b/2DTM_postprocess_tool/src/tm_post/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/2DTM_postprocess_tool/src/tm_post/compare_starfiles.py b/2DTM_postprocess_tool/src/tm_post/compare_starfiles.py new file mode 100644 index 000000000..e091ed9a4 --- /dev/null +++ b/2DTM_postprocess_tool/src/tm_post/compare_starfiles.py @@ -0,0 +1,70 @@ +import pandas as pd +from scipy.spatial import cKDTree +import re +from tm_post.starfile import load_particle_starfile +from tm_post.geometry import return_euler_err + +def extract_match_key(filename, pattern): + match = re.search(pattern, filename) + return match.group(1) if match else None + +def compare_starfiles_for_matched_peaks( + starfile_a: str, + starfile_b: str, + d_xy_cutoff: float = 10.0, # in Å + euler_err_cutoff: float = 5.0, # in degrees + pattern=r"mc2_[12]x_(.*?frames)" +) -> pd.DataFrame: + """ + Compare two STAR DataFrames (df_a and df_b) and return matched peaks. + Accounts for pixel size differences between the two searches. + """ + # Load the starfile + df_a = load_particle_starfile(starfile_a) + df_b = load_particle_starfile(starfile_b) + df_a["match_key"] = df_a["ORIGINAL_IMAGE_FILENAME"].apply(lambda x: extract_match_key(x, pattern)) + df_b["match_key"] = df_b["ORIGINAL_IMAGE_FILENAME"].apply(lambda x: extract_match_key(x, pattern)) + + # Find overlapped images + common_keys = set(df_a["match_key"]) & set(df_b["match_key"]) + print(f"[INFO] Found {len(common_keys)} overlapping images between the two searches.") + + matched_idx_b = [] + # Iterate through each common key + for key in common_keys: + peaks_a = df_a[df_a["match_key"] == key] + peaks_b = df_b[df_b["match_key"] == key] + + if peaks_a.empty or peaks_b.empty: + continue + + coords_a = peaks_a[["ORIGX", "ORIGY"]].values + coords_b = peaks_b[["ORIGX", "ORIGY"]].values + + tree_b = cKDTree(coords_b) + dists, idxs_b = tree_b.query(coords_a) + + for i, (dist, j) in enumerate(zip(dists, idxs_b)): + if dist >= d_xy_cutoff: + continue + + row_a = peaks_a.iloc[i] + row_b = peaks_b.iloc[j] + + err = return_euler_err( + gt_psi=row_a["PSI"], gt_theta=row_a["THETA"], gt_phi=row_a["PHI"], + tm_psi=row_b["PSI"], tm_theta=row_b["THETA"], tm_phi=row_b["PHI"] + ) + + if err < euler_err_cutoff: + idx_b = peaks_b.index[j] + matched_idx_b.append(idx_b) + + # Annotate matched rows in df_b + df_b_matched = df_b.loc[matched_idx_b].copy() + df_b_matched.drop(columns=["match_key"], inplace=True) + print(f"[INFO] Found {len(df_b_matched)} matched peaks.") + + return df_b_matched + + \ No newline at end of file diff --git a/2DTM_postprocess_tool/src/tm_post/database.py b/2DTM_postprocess_tool/src/tm_post/database.py new file mode 100644 index 000000000..a1a6e68cd --- /dev/null +++ b/2DTM_postprocess_tool/src/tm_post/database.py @@ -0,0 +1,148 @@ +import sqlite3 +import pandas as pd +from tm_post.image_data import TMImage + +def get_info_from_cistem_database(db_file, tm_job_id, ctf_job_id, requested_output_names=None): + # Table names: CTF, image info, TM job + table_ctf = 'ESTIMATED_CTF_PARAMETERS' + table_info = "IMAGE_ASSETS" + table_tm = 'TEMPLATE_MATCH_LIST' + + # Connect to the database + conn = sqlite3.connect(db_file) + cursor = conn.cursor() + + # Load all three tables into dataframes + def load_table(query): + cursor.execute(query) + rows = cursor.fetchall() + columns = [desc[0] for desc in cursor.description] + return pd.DataFrame(rows, columns=columns) + + df_ctf = load_table(f"SELECT * FROM {table_ctf}") + df_info = load_table(f"SELECT * FROM {table_info}") + df_tm = load_table(f"SELECT * FROM {table_tm}") + + # Filter CTF data + df_ctf = df_ctf[df_ctf['CTF_ESTIMATION_JOB_ID'] == ctf_job_id] + + # Filter TM data + df_tm = df_tm[df_tm['TEMPLATE_MATCH_JOB_ID'] == tm_job_id] + + # Close the connection + conn.close() + + # Get relevant image IDs + image_ids = df_tm.IMAGE_ASSET_ID.values + + # Always needed + image_list = [df_info[df_info.IMAGE_ASSET_ID == image_id].FILENAME.values[0] for image_id in image_ids] + + # Define mapping of column name to list + all_output_columns = { + "MIP_OUTPUT_FILE": [], + "SCALED_MIP_OUTPUT_FILE": [], + "PSI_OUTPUT_FILE": [], + "THETA_OUTPUT_FILE": [], + "PHI_OUTPUT_FILE": [], + "DEFOCUS_OUTPUT_FILE": [], + "AVG_OUTPUT_FILE": [], + "STD_OUTPUT_FILE": [] + } + + # Default is return all output columns + if requested_output_names is None: + requested_output_names = all_output_columns + elif isinstance(requested_output_names, str): + requested_output_names = [requested_output_names] + + # Create output lists + output_lists = {col: [] for col in requested_output_names} + + for image_id in image_ids: + row_tm = df_tm[df_tm.IMAGE_ASSET_ID == image_id] + + # Output files + for col in requested_output_names: + if col not in output_lists: + raise ValueError(f"Column '{col}' is not recognized.") + if col not in row_tm.columns: + raise ValueError(f"Column '{col}' not found in TM database table.") + output_lists[col].append(row_tm[col].values[0]) + + # Assemble result + result = { + "image_list": image_list, + "image_ids": image_ids, + "df_ctf": df_ctf, + "df_info" : df_info + } + + for col in requested_output_names: + result[col] = output_lists[col] + + return result + +def load_tm_images_from_db(db_file, tm_job_id, ctf_job_id):# -> list[TMImage]: + # Table names: CTF, image info, TM job + table_ctf = 'ESTIMATED_CTF_PARAMETERS' + table_info = "IMAGE_ASSETS" + table_tm = 'TEMPLATE_MATCH_LIST' + + # Connect to the database + conn = sqlite3.connect(db_file) + cursor = conn.cursor() + + # Load all three tables into dataframes + def load_table(query): + cursor.execute(query) + rows = cursor.fetchall() + columns = [desc[0] for desc in cursor.description] + return pd.DataFrame(rows, columns=columns) + + df_ctf = load_table(f"SELECT * FROM {table_ctf}") + df_info = load_table(f"SELECT * FROM {table_info}") + df_tm = load_table(f"SELECT * FROM {table_tm}") + + # Filter CTF data + df_ctf = df_ctf[df_ctf['CTF_ESTIMATION_JOB_ID'] == ctf_job_id] + + # Filter TM data + df_tm = df_tm[df_tm['TEMPLATE_MATCH_JOB_ID'] == tm_job_id] + + # Close the connection + conn.close() + + # Get relevant image IDs + image_ids = df_tm.IMAGE_ASSET_ID.values + + # Instead of returning lists, build TMImage objects + images = [] + + for image_id in image_ids: + filename = df_info[df_info.IMAGE_ASSET_ID == image_id].FILENAME.values[0] + row_tm = df_tm[df_tm.IMAGE_ASSET_ID == image_id] + row_ctf = df_ctf[df_ctf.IMAGE_ASSET_ID == image_id] + + image = TMImage( + image_id=image_id, + filename=filename, + pixel_size=row_tm["USED_PIXEL_SIZE"].values[0], + psi_file=row_tm["PSI_OUTPUT_FILE"].values[0], + theta_file=row_tm["THETA_OUTPUT_FILE"].values[0], + phi_file=row_tm["PHI_OUTPUT_FILE"].values[0], + snr_file=row_tm["MIP_OUTPUT_FILE"].values[0], + zscore_file=row_tm["SCALED_MIP_OUTPUT_FILE"].values[0], + defocus_file=row_tm["DEFOCUS_OUTPUT_FILE"].values[0], + avg_file=row_tm["AVG_OUTPUT_FILE"].values[0], + sd_file=row_tm["STD_OUTPUT_FILE"].values[0], + defocus1=row_ctf["DEFOCUS1"].values[0], + defocus2=row_ctf["DEFOCUS2"].values[0], + defocus_angle=row_ctf["DEFOCUS_ANGLE"].values[0], + amp_contrast=row_ctf["AMPLITUDE_CONTRAST"].values[0], + voltage=row_ctf["VOLTAGE"].values[0], + cs=row_ctf["SPHERICAL_ABERRATION"].values[0], + ) + images.append(image) + + return images, df_ctf, df_info \ No newline at end of file diff --git a/2DTM_postprocess_tool/src/tm_post/extract.py b/2DTM_postprocess_tool/src/tm_post/extract.py new file mode 100644 index 000000000..5fa5b7681 --- /dev/null +++ b/2DTM_postprocess_tool/src/tm_post/extract.py @@ -0,0 +1,136 @@ + +from tm_post.peak import Peak +from tm_post.image_data import TMImage +from skimage.feature import peak_local_max +import pandas as pd +from tm_post.statistics import calculate_2dtm_pval +from tm_post.mrcfile import read_mrc_file +from tm_post.starfile import convert_peaks_to_star_df +import concurrent.futures +from tqdm import tqdm + +def return_peaks_for_image(image: TMImage, metric_cutoff, local_max_filter="zscore", metric="pval", min_radius=10, exclude_borders=35, q=3):# -> list[Peak]: + """Generate peak information for a given image in a database.""" + # Read all maps + snr_image = read_mrc_file(image.snr_file) + zscore_image = read_mrc_file(image.zscore_file) + psi_image = read_mrc_file(image.psi_file) + theta_image = read_mrc_file(image.theta_file) + phi_image = read_mrc_file(image.phi_file) + defocus_image = read_mrc_file(image.defocus_file) + avg_image = read_mrc_file(image.avg_file) + sd_image = read_mrc_file(image.sd_file) + + if local_max_filter == "zscore": + peaks_coordinates = peak_local_max(zscore_image, min_distance=min_radius, exclude_border=exclude_borders, threshold_abs=0.0) + elif local_max_filter == "snr": + peaks_coordinates = peak_local_max(snr_image, min_distance=min_radius, exclude_border=exclude_borders, threshold_abs=0.0) + + # Collect raw values from detected peaks + peak_data = [] + for (y,x) in peaks_coordinates: + peak_data.append({ + "x_pixel": x, + "y_pixel": y, + "snr": snr_image[y, x], + "zscore": zscore_image[y, x], + "psi": psi_image[y, x], + "theta": theta_image[y, x], + "phi": phi_image[y, x], + "delta_defocus": defocus_image[y, x], + "avg": avg_image[y, x], + "sd": sd_image[y, x], + }) + + # Compute p-values + df_peaks = pd.DataFrame(peak_data) + df_peaks["pval"] = calculate_2dtm_pval(df_peaks["zscore"].values, df_peaks["snr"].values, q=q) + + # Filter and create Peak objects + filtered_peaks = [] + for _, row in df_peaks.iterrows(): + #if row["avg"] > avg_cutoff and row["snr"] < snr_cutoff: + if (metric == "pval" and row["pval"] >= metric_cutoff) or \ + (metric == "zscore" and row["zscore"] >= metric_cutoff) or \ + (metric == "snr" and row["snr"] >= metric_cutoff): + filtered_peaks.append( + Peak( + image_id=image.image_id, + image_name=image.filename, + x=row["x_pixel"], + y=row["y_pixel"], + delta_defocus=row["delta_defocus"], + psi=row["psi"], + theta=row["theta"], + phi=row["phi"], + snr=row["snr"], + zscore=row["zscore"], + pval=row["pval"], + avg=row["avg"], + sd=row["sd"], + ) + ) + + # Sort peaks by selected metric in descending order + filtered_peaks = sorted(filtered_peaks, key=lambda p: getattr(p, metric), reverse=True) + + print(f"[INFO] Extracted {len(filtered_peaks)} peaks from image: {image.filename}") + return filtered_peaks + +def extract_particles_from_2dtm_search( + tm_images, + local_max_filter, + df_ctf, + df_info, + ctf_job_id, + metric = "pval", + metric_cutoff = 8.0, + #avg_cutoff = 0.0, + #snr_cutoff = 9.0, + pixel_size = 1.0, + min_radius = 10, + exclude_borders = 35, + max_threads = 4, + q = 3, + ): + """ + Extract particles (peaks) from a list of TMImage objects in parallel. + + Returns a full STAR-format DataFrame of all particles across all images. + """ + all_particles = [] + def process_image(image: TMImage): + # extract Peak objects for this image + peaks = return_peaks_for_image( + image=image, + #avg_cutoff=avg_cutoff, + #snr_cutoff=snr_cutoff, + local_max_filter=local_max_filter, + metric_cutoff=metric_cutoff, + metric=metric, + min_radius=min_radius, + exclude_borders=exclude_borders, + q=q + ) + print(f"[INFO] Found {len(peaks)} peaks in image {image.filename}") + + # convert to STAR-format DataFrame + return convert_peaks_to_star_df( + peaks=peaks, + image_id=image.image_id, + df_ctf=df_ctf, + df_info=df_info, + ctf_job_id=ctf_job_id, + pixel_size=pixel_size, + multiply_pixel_size=True, + metric=metric + ) + + all_dataframes = [] + with concurrent.futures.ThreadPoolExecutor(max_workers=max_threads) as executor: + futures = [executor.submit(process_image, img) for img in tm_images] + for future in tqdm(concurrent.futures.as_completed(futures), total=len(futures), desc="Extracting particles"): + df = future.result() + all_dataframes.append(df) + + return pd.concat(all_dataframes, ignore_index=True) \ No newline at end of file diff --git a/2DTM_postprocess_tool/src/tm_post/filters.py b/2DTM_postprocess_tool/src/tm_post/filters.py new file mode 100644 index 000000000..2eec43159 --- /dev/null +++ b/2DTM_postprocess_tool/src/tm_post/filters.py @@ -0,0 +1,131 @@ +import numpy as np +import pandas as pd +import operator + +from tm_post.geodesic import calculate_all_geodesic_means + + +def apply_image_thickness_filter(peaks, df_ctf, df_info, cutoff_lb, cutoff_ub): + # Get filenames of images within thickness cutoff + df_ctf_filtered = df_ctf[ + (df_ctf['SAMPLE_THICKNESS'] > cutoff_lb) & + (df_ctf['SAMPLE_THICKNESS'] < cutoff_ub) + ] + + filenames = df_info[df_info['IMAGE_ASSET_ID'].isin(df_ctf_filtered['IMAGE_ASSET_ID'])]['FILENAME'] + filenames = [f"'{filename}'" for filename in filenames] # keep quotes if needed + + # Create a boolean mask + mask = peaks['ORIGINAL_IMAGE_FILENAME'].isin(filenames) + + return mask + +def apply_ctf_score_filter(peaks, df_ctf, df_info, cutoff_lb, cutoff_ub): + df_ctf_filtered = df_ctf[ + (df_ctf['SCORE'] > cutoff_lb) & + (df_ctf['SCORE'] < cutoff_ub) + ] + filenames = df_info[df_info['IMAGE_ASSET_ID'].isin(df_ctf_filtered['IMAGE_ASSET_ID'])]['FILENAME'] + filenames = [f"'{filename}'" for filename in filenames] # keep quotes if needed + # Create a boolean mask + mask = peaks['ORIGINAL_IMAGE_FILENAME'].isin(filenames) + return mask + +def apply_angular_invariance_filter(df, mean_geodesic_array, method='quantile', threshold=0.95): + """ + Create a lookup function for image thickness based on ORIGINAL_IMAGE_FILENAME. + Returns a function that takes a filename string and returns the sample thickness. + """ + if method == 'quantile': + cutoff = np.nanquantile(mean_geodesic_array, threshold) + elif method == 'cutoff': + cutoff = threshold + else: + raise ValueError("method must be 'quantile' or 'cutoff'") + + #df_filtered = df[keep_mask].reset_index(drop=True) + return mean_geodesic_array < cutoff, cutoff + + +def apply_filter( + df, + image_list, + psi_list, + theta_list, + phi_list, + pixel_size, + df_ctf, + df_info, + avg_cutoff_lb=None, + sd_cutoff_ub=None, + pval_cutoff_lb=None, + snr_cutoff_lb=None, + snr_cutoff_ub=None, + ctf_fitting_score_lb=None, + ctf_fitting_score_ub=None, + filter_by_image_thickness=True, + thickness_lb=None, + thickness_ub=None, + filter_by_angular_invariance=False, + geodesic_r=4, + geodesic_threads=8, + geodesic_method='quantile', # or 'cutoff' + geodesic_threshold=0.8 # quantile (0.8) or distance cutoff (e.g., 0.3) +): + """ + Apply multi-criteria filtering to template matching results. + """ + df_filtered = df.copy() + df_record = df_filtered[['ORIGINAL_IMAGE_FILENAME','ORIGX','ORIGY','AVG','SD','PVALUE','ZSCORE','SNR']].copy() + + # Initialize to all True + current_mask = pd.Series(True, index=df_filtered.index) + geodesic_means = None + + # Define standard numeric filters (col, value, operator) + value_filters = [ + ('AVG', avg_cutoff_lb, operator.gt), + ('SD', sd_cutoff_ub, operator.lt), + ('SNR', snr_cutoff_lb, operator.gt), + ('SNR', snr_cutoff_ub, operator.lt), + ('PVALUE', pval_cutoff_lb, operator.gt), + ] + + # Apply scalar filters + for name, cutoff, op in value_filters: + if cutoff is not None: + mask = op(df_filtered[name], cutoff) + current_mask &= mask + print(f"[INFO] Filter `{name} {op.__name__} {cutoff}`: {mask.sum()} particles retained") + + + if filter_by_image_thickness and thickness_lb is not None and thickness_ub is not None: + thickness_mask = apply_image_thickness_filter(df_filtered, df_ctf, df_info, thickness_lb, thickness_ub) + current_mask &= thickness_mask + print(f"[INFO] Thickness filter [{thickness_lb}, {thickness_ub}]: {current_mask.sum()} particles retained") + + if ctf_fitting_score_lb is not None and ctf_fitting_score_ub is not None: + ctf_mask = apply_ctf_score_filter(df_filtered, df_ctf, df_info, ctf_fitting_score_lb, ctf_fitting_score_ub) + current_mask &= ctf_mask + print(f"[INFO] CTF SCORE filter [{ctf_fitting_score_lb}, {ctf_fitting_score_ub}]: {current_mask.sum()} particles retained") + + # Apply angular invariance filtering + if filter_by_angular_invariance: + print(f"[INFO] Calculating angular variance...") + geodesic_means = calculate_all_geodesic_means( + df_filtered, image_list, psi_list, theta_list, phi_list, + pixel_size, r=geodesic_r, threads=geodesic_threads + ) + # Keep only thickness-filtered geodesic values + df_record['mean_geodesic_distance'] = geodesic_means + + angular_mask, cutoff_val = apply_angular_invariance_filter( + df_filtered, geodesic_means, method=geodesic_method, threshold=geodesic_threshold + ) + current_mask &= angular_mask + print(f"[INFO] Geodesic filter ({geodesic_method} ≤ {cutoff_val:.3f}): {current_mask.sum()} particles retained") + + # Final mask application + df_filtered = df_filtered[current_mask].reset_index(drop=True) + print(f"[INFO] Final filtered particles: {len(df_filtered)}") + return df_filtered, df_record diff --git a/2DTM_postprocess_tool/src/tm_post/geodesic.py b/2DTM_postprocess_tool/src/tm_post/geodesic.py new file mode 100644 index 000000000..a48ee3252 --- /dev/null +++ b/2DTM_postprocess_tool/src/tm_post/geodesic.py @@ -0,0 +1,92 @@ +import numpy as np +import mrcfile +from tqdm import tqdm +import concurrent.futures +from tm_post.geometry import geodesic_distance, euler_to_rotation + +def get_local_patch(file_name, x, y, r): + """Extract a square patch around (x, y) from an MRC file.""" + with mrcfile.open(file_name) as mrc: + data = mrc.data[0] # Assuming it's 3D with first slice of interest + return data[y - r:y + r, x - r:x + r] + +def rotation_matrix_patch(psi_patch, theta_patch, phi_patch): + """Convert Euler angle patches to rotation matrix patches.""" + r = psi_patch.shape[0] // 2 + rot_patch = np.empty((2*r, 2*r), dtype=object) # Store Rotation objects + for i in range(2*r): + for j in range(2*r): + rot_patch[i, j] = euler_to_rotation( + psi_patch[i, j], theta_patch[i, j], phi_patch[i, j] + ) + return rot_patch + +def compute_geodesic_distances(rot_patch, center_coord): + """Compute geodesic distances from center rotation to all others.""" + r = rot_patch.shape[0] // 2 + ref_rot = rot_patch[r, r] + distances = [] + for i in range(2*r): + for j in range(2*r): + if i != r or j != r: # exclude center + dist = geodesic_distance(ref_rot, rot_patch[i, j]) + distances.append(dist) + return distances + +def calculate_particle_geodesic_distance(df, image_list, psi_list, theta_list, phi_list, peak_number, pixel_size, r=10): + """Top-level function to calculate geodesic distances near a particle.""" + # Get particle info + row = df.iloc[peak_number-1] + x, y = row['ORIGX'] / pixel_size, row['ORIGY'] / pixel_size + x, y = int(round(x,1)), int(round(y,1)) + image_name = row['ORIGINAL_IMAGE_FILENAME'].strip("'") + + image_idx = image_list.index(image_name) + psi_file, theta_file, phi_file = psi_list[image_idx], theta_list[image_idx], phi_list[image_idx] + + # Load patches + psi_patch = get_local_patch(psi_file, x, y, r) + theta_patch = get_local_patch(theta_file, x, y, r) + phi_patch = get_local_patch(phi_file, x, y, r) + + # Convert to rotation matrices and compute distances + rot_patch = rotation_matrix_patch(psi_patch, theta_patch, phi_patch) + distances = compute_geodesic_distances(rot_patch, (r, r)) + + return distances + + +def calculate_mean_geodesic_for_row(args): + """Wrapper to calculate mean geodesic distance for one row.""" + df, image_list, psi_list, theta_list, phi_list, pixel_size, r, row_idx = args + try: + distances = calculate_particle_geodesic_distance( + df, image_list, psi_list, theta_list, phi_list, + peak_number=row_idx + 1, # your function is 1-based + pixel_size=pixel_size, + r=r + ) + return np.mean(distances) + except Exception as e: + print(f"[Warning] Skipped row {row_idx} due to error: {e}") + return np.nan + +def calculate_all_geodesic_means(df, image_list, psi_list, theta_list, phi_list, pixel_size, r=10, threads=4): + """Compute mean geodesic distance for each row in df in parallel, with progress bar.""" + args_list = [ + (df, image_list, psi_list, theta_list, phi_list, pixel_size, r, i) + for i in range(len(df)) + ] + means = [None] * len(args_list) + + with concurrent.futures.ThreadPoolExecutor(max_workers=threads) as executor: + futures = {executor.submit(calculate_mean_geodesic_for_row, args): idx for idx, args in enumerate(args_list)} + for future in tqdm(concurrent.futures.as_completed(futures), total=len(futures), desc="Calculating geodesic means"): + idx = futures[future] + try: + means[idx] = future.result() + except Exception as e: + print(f"[Warning] Error in row {idx}: {e}") + means[idx] = np.nan + + return np.array(means) \ No newline at end of file diff --git a/2DTM_postprocess_tool/src/tm_post/geometry.py b/2DTM_postprocess_tool/src/tm_post/geometry.py new file mode 100644 index 000000000..a665ab2fd --- /dev/null +++ b/2DTM_postprocess_tool/src/tm_post/geometry.py @@ -0,0 +1,60 @@ +import numpy as np +from scipy.spatial.transform import Rotation as R + +def euler_to_rotation(psi, theta, phi, degrees=True): + """ + Convert Euler angles (psi, theta, phi) to a Rotation object using the ZYZ convention. + The rotations are applied in the order: + 1. Rotation about the Z-axis by psi, + 2. Rotation about the Y-axis by theta, + 3. Rotation about the Z-axis by phi. + + This convention aligns with the angular annotations used in RELION and cisTEM. + """ + return R.from_euler('ZYZ', [psi, theta, phi], degrees=degrees) + +def euler_to_matrix(psi, theta, phi, degrees=True): + return R.from_euler('ZYZ', [psi, theta, phi], degrees=degrees).as_matrix() + +def geodesic_distance(ref_rot, pixel_rot): + """ + Compute the geodesic distance (in radians) between two rotations. + + Parameters: + - ref_rot: the reference Rotation object. + - pixel_rot: the Rotation object of the pixel to compare. + + The geodesic distance is the magnitude of the rotation vector + corresponding to the relative rotation between ref_rot and pixel_rot. + """ + # Compute the relative rotation from reference to pixel + relative_rot = ref_rot.inv() * pixel_rot + # The magnitude of the rotation vector represents the geodesic distance + angle = np.linalg.norm(relative_rot.as_rotvec()) + return angle + + +def Rz(x_azimu): + x_azimu = x_azimu*np.pi/180 + return np.array([[+np.cos(x_azimu), -np.sin(x_azimu), 0], [+np.sin(x_azimu), +np.cos(x_azimu), 0], [0, 0, 1]]) + +def Ry(x_polar): + x_polar = x_polar*np.pi/180 + return np.array([[+np.cos(x_polar), 0, +np.sin(x_polar)], [0, 1, 0], [-np.sin(x_polar), 0, +np.cos(x_polar)]]) + +def return_euler_err(gt_psi,gt_theta,gt_phi, tm_psi, tm_theta,tm_phi): + n_gamma_z = 1024 + gamma_z_ = np.linspace(0,2*np.pi,n_gamma_z+1) + gamma_z_ = np.transpose(gamma_z_[0:n_gamma_z]) + ring_k_c_0_ = np.cos(gamma_z_) + ring_k_c_1_ = np.sin(gamma_z_) + ring_k_c_2_ = np.zeros(n_gamma_z) + ring_k_c_3z__ = np.array([ring_k_c_0_,ring_k_c_1_,ring_k_c_2_]) + tmp_ring_est_k_c_3z__ = np.matmul(Rz(tm_phi),Ry(tm_theta)) + tmp_ring_est_k_c_3z__ = np.matmul(tmp_ring_est_k_c_3z__, Rz(tm_psi)) + tmp_ring_est_k_c_3z__ = np.matmul(tmp_ring_est_k_c_3z__, ring_k_c_3z__) + tmp_ring_tru_k_c_3z__ = np.matmul(Rz(gt_phi),Ry(gt_theta)) + tmp_ring_tru_k_c_3z__ = np.matmul(tmp_ring_tru_k_c_3z__, Rz(gt_psi)) + tmp_ring_tru_k_c_3z__ = np.matmul(tmp_ring_tru_k_c_3z__, ring_k_c_3z__) + tmp_ring_l2 = np.sqrt(np.sum((tmp_ring_est_k_c_3z__ - tmp_ring_tru_k_c_3z__)**2)*2*np.pi/max(1,n_gamma_z)) + return tmp_ring_l2 \ No newline at end of file diff --git a/2DTM_postprocess_tool/src/tm_post/image_data.py b/2DTM_postprocess_tool/src/tm_post/image_data.py new file mode 100644 index 000000000..25705ebc3 --- /dev/null +++ b/2DTM_postprocess_tool/src/tm_post/image_data.py @@ -0,0 +1,25 @@ +from dataclasses import dataclass + +@dataclass +class TMImage: + image_id: int + filename: str + + snr_file: str + zscore_file: str + avg_file: str + sd_file: str + + psi_file: str + theta_file: str + phi_file: str + + defocus_file: str + pixel_size: float + defocus1: float + defocus2: float + defocus_angle: float + + amp_contrast: float + voltage: float + cs: float \ No newline at end of file diff --git a/2DTM_postprocess_tool/src/tm_post/mrcfile.py b/2DTM_postprocess_tool/src/tm_post/mrcfile.py new file mode 100644 index 000000000..955f1642f --- /dev/null +++ b/2DTM_postprocess_tool/src/tm_post/mrcfile.py @@ -0,0 +1,8 @@ +import mrcfile +import numpy as np + +def read_mrc_file(filename): + with mrcfile.open(filename) as mrc: + df = np.squeeze(mrc.data) + + return df \ No newline at end of file diff --git a/2DTM_postprocess_tool/src/tm_post/peak.py b/2DTM_postprocess_tool/src/tm_post/peak.py new file mode 100644 index 000000000..b9fedc446 --- /dev/null +++ b/2DTM_postprocess_tool/src/tm_post/peak.py @@ -0,0 +1,50 @@ +from dataclasses import dataclass +import numpy as np +from scipy.spatial.transform import Rotation as R + +@dataclass +class Peak: + image_id: int + image_name: str + x: float # in pixel + y: float # in pixel + delta_defocus: float + psi: float + theta: float + phi: float + snr: float + zscore: float + pval: float + avg: float + sd: float + + #def get_peak_coordinates(self): + # return np.array([self.x * self.pixel_size, self.y * self.pixel_size]) + + def convert_to_rotation_matrix(self): + # Create rotation matrix from Euler angles (psi, theta, phi) + return R.from_euler('ZYZ', [self.psi, self.theta, self.phi], degrees=True) + + def convert_to_starfile_row(self): + return [ + "", # empty column + round(self.psi, 1), + round(self.theta, 1), + round(self.phi, 1), + round(self.defocus1, 1), + round(self.defocus2, 1), + 0.0, # ANGAST placeholder + round(self.pixel_size, 3), + 200.0, # microscope voltage + 2.7, # Cs + 0.1, # Amp contrast + 0.0, 0.0, # beam tilt X/Y + 0.0, 0.0, # image shift X/Y + f"'{self.image_name}'", + round(self.x, 2), + round(self.y, 2), + round(self.pval, 2), + round(self.zscore, 2), + round(self.snr, 2) + ] + \ No newline at end of file diff --git a/2DTM_postprocess_tool/src/tm_post/starfile.py b/2DTM_postprocess_tool/src/tm_post/starfile.py new file mode 100644 index 000000000..afea1f916 --- /dev/null +++ b/2DTM_postprocess_tool/src/tm_post/starfile.py @@ -0,0 +1,166 @@ +import pandas as pd +from tm_post.peak import Peak +from tm_post.image_data import TMImage +from pathlib import Path + +def read_tm_package_starfile_header(): #-> list[str]: + """ + Reads the default STAR file header from sample_data/header.star. + + Returns: + A list of header lines (each ending in '\n'). + """ + header_path = Path(__file__).resolve().parents[2] / "sample_data" / "header.star" + with open(header_path, "r") as f: + lines = f.readlines() + return lines + +def add_star_dummy_column(df: pd.DataFrame) -> pd.DataFrame: + """Insert a dummy column named '#' as the first column, filled with empty strings.""" + df_out = df.copy() + df_out.insert(0, "#", "") + return df_out + +def extract_header_lines(file_path): + """Extract real header lines (e.g., lines starting with '#' and containing PSI, etc.).""" + header_keywords = ["PSI", "THETA", "PHI", "DF1", "SCORE"] + header_lines = [] + data_lines = [] + + with open(file_path, 'r') as f: + for line in f: + if line.lstrip().startswith('#') and all(key in line.upper() for key in header_keywords): + header_lines.append(line) + else: + data_lines.append(line) + + return header_lines, data_lines + +def write_starfile_with_headers(filepath, header_lines, df): + """ + Write a STAR file with preformatted header lines and a data DataFrame that already includes the "#" column. + + Parameters: + - filepath: path to the output file + - header_lines: list of strings starting with "#", already including newline + - df: DataFrame with the "#" column already included as the first column + """ + with open(filepath, 'w') as f: + f.writelines(header_lines) + f.write("\n") # Add a blank line between header and data + df.to_csv(f, sep="\t", index=False, header=False) + +def find_data_header_lines(file_path): + """Find line numbers of header lines that contain particle column names.""" + header_keywords = ["PSI", "THETA", "PHI", "DF1", "SCORE"] # feel free to expand + header_lines = [] + with open(file_path, 'r') as f: + for idx, line in enumerate(f, start=1): + if line.lstrip().startswith('#') and all(key in line.upper() for key in header_keywords): + header_lines.append(idx) + return header_lines + +def load_particle_starfile(file_path): + """Load particle data from a starfile.""" + # Read in starfile and find header lines + header_lines = find_data_header_lines(file_path) + if not header_lines: + raise ValueError("No header lines found in the starfile.") + # Read the starfile, skipping the header lines + df = pd.read_csv(file_path, delim_whitespace=True, skiprows=header_lines[0], header=None) + # Set column names based on column indices + if df.shape[1]==23: + df.columns = ["PSI", "THETA", "PHI", "DF1", "DF2", "ANGAST", + "SCORE", "PSIZE", "VOLT", "Cs", "AmpC", + "BTILTX", "BTILTY", "ISHFTX", "ISHFTY", + "ORIGINAL_IMAGE_FILENAME", "ORIGX", "ORIGY", + "PVALUE", "ZSCORE", "SNR", "AVG", "SD" + ] + elif df.shape[1]==18: + df.columns = ["PSI", "THETA", "PHI", "DF1", "DF2", "ANGAST", + "SCORE", "PSIZE", "VOLT", "Cs", "AmpC", + "BTILTX", "BTILTY", "ISHFTX", "ISHFTY", + "ORIGINAL_IMAGE_FILENAME", "ORIGX", "ORIGY" + ] + elif df.shape[1]==24: # from binary + df.columns = ["POS", "PSI", "THETA", "PHI", "SHX", "SHY", "DF1", + "DF2", "ANGAST", "PSHIFT", "STAT", "OCC", + "LogP", "SIGMA", "SCORE", "PSIZE", + "VOLT", "Cs", "AmpC", "BTILTX", "BTILTY", + "ISHFTX", "ISHFTY", "SUBSET", + ] + elif df.shape[1]==29: # from simulator + df.columns = ["POS", "PSI", "THETA", "PHI", "SHX", "SHY", + "DF1", "DF2", "ANGAST", "PSHIFT", "OCC", + "LogP", "SIGMA", "SCORE", "CHANGE", "PSIZE", + "VOLT", "Cs", "AmpC", "BTILTX", "BTILTY", + "ISHFTX", "ISHFTY", "2DCLS", "TGRP", "PaGRP", + "SUBSET", "PREEXP", "TOTEXP"] + return df + + + +def convert_peaks_to_star_df(peaks,image_id,df_ctf,df_info,ctf_job_id,pixel_size,metric="pval",multiply_pixel_size= False +): + """ + Convert a list of Peak objects to a STAR-format DataFrame. + Parameters: + - peaks: list of Peak objects + - image_id: ID of the image + - df_ctf: DataFrame containing CTF information + - df_info: DataFrame containing image information + - ctf_job_id: CTF job ID + - pixel_size: pixel size in Angstroms + - multiply_pixel_size: whether to multiply x and y coordinates by pixel size + Returns: + - df_out: DataFrame in STAR format + """ + if not peaks: + return pd.DataFrame() + + # Get image metadata + row_ctf = df_ctf[(df_ctf.CTF_ESTIMATION_JOB_ID == ctf_job_id) & (df_ctf.IMAGE_ASSET_ID == image_id)] + row_info = df_info[df_info.IMAGE_ASSET_ID == image_id] + + defocus1 = row_ctf['DEFOCUS1'].values[0] + defocus2 = row_ctf['DEFOCUS2'].values[0] + defocus_angle = row_ctf['DEFOCUS_ANGLE'].values[0] + filename = row_info['FILENAME'].values[0] + cs = row_info['SPHERICAL_ABERRATION'].values[0] + voltage = row_info['VOLTAGE'].values[0] + amp_contrast = row_ctf['AMPLITUDE_CONTRAST'].values[0] + # Build STAR-format DataFrame + df_out = pd.DataFrame({ + "#": ["" for _ in peaks], + "PSI": [round(p.psi, 1) for p in peaks], + "THETA": [round(p.theta, 1) for p in peaks], + "PHI": [round(p.phi, 1) for p in peaks], + "DF1": [round(p.delta_defocus + defocus1,1) for p in peaks], + "DF2": [round(p.delta_defocus + defocus2,1) for p in peaks], + "ANGAST": [round(defocus_angle,1) for _ in peaks], + "SCORE": [ + round( + p.pval if metric == "pval" + else p.zscore if metric == "zscore" + else p.snr if metric == "snr" + else float("nan"), 2 + ) for p in peaks + ], + "PSIZE": [pixel_size for _ in peaks], + "VOLT": [round(voltage,1) for _ in peaks], + "Cs": [round(cs,1) for _ in peaks], + "AmpC": [round(amp_contrast,3) for _ in peaks], # adjust as needed + "BTILTX": [0.0 for _ in peaks], + "BTILTY": [0.0 for _ in peaks], + "ISHFTX": [0.0 for _ in peaks], + "ISHFTY": [0.0 for _ in peaks], + "ORIGINAL_IMAGE_FILENAME": [f"'{filename}'" for _ in peaks], + "ORIGX": [round(p.x * pixel_size if multiply_pixel_size else p.x,2) for p in peaks], + "ORIGY": [round(p.y * pixel_size if multiply_pixel_size else p.y,2) for p in peaks], + "PVALUE": [round(p.pval, 2) for p in peaks], + "ZSCORE": [round(p.zscore, 2) for p in peaks], + "SNR": [round(p.snr, 2) for p in peaks], + "AVG": [round(p.avg, 2) for p in peaks], + "SD": [round(p.sd, 2) for p in peaks] + }) + return df_out diff --git a/2DTM_postprocess_tool/src/tm_post/statistics.py b/2DTM_postprocess_tool/src/tm_post/statistics.py new file mode 100644 index 000000000..c25233f0a --- /dev/null +++ b/2DTM_postprocess_tool/src/tm_post/statistics.py @@ -0,0 +1,117 @@ +import numpy as np +import math +from scipy.special import erfinv + +def calculate_probit(x_): + # calculate rank and remap to quantile of standard gaussian + rank_x_ = np.argsort(np.argsort(x_))+1 # use argsort twice to ensure the ranking is correct + rank_x_ = (rank_x_ - 0.5) / max(1, len(x_)) + pro_x_ = np.sqrt(2) * erfinv(2 * rank_x_ - 1.0) + return rank_x_, pro_x_ + +def estimate_anisotropic_gaussian_for_probit(pro_x1_, pro_x2_): + pro_lim_ = [-4.5, 4.5] + # define anisotropic-gaussian + n_pro_x = 128 + n_pro_y = n_pro_x+1 + pro_x_ = np.linspace(min(pro_lim_), max(pro_lim_), n_pro_x) + pro_dx = np.mean(np.diff(pro_x_)) + + pro_y_ = np.linspace(min(pro_lim_), max(pro_lim_), n_pro_y) + pro_dy = np.mean(np.diff(pro_y_)) + + pro_x__, pro_y__ = np.meshgrid(pro_x_,pro_y_) + + n_r = len(pro_x1_) + tmp_ = np.array([pro_x1_, pro_x2_]) + Cinv__ = np.matmul(tmp_, tmp_.T)/n_r + Dinv_,Uinv__ = np.linalg.eig(Cinv__) # eigenvalues not necessarily ordered + return Dinv_, Uinv__ #C__, g__, pro_x__,pro_y__, pro_dx, pro_dy + + +def calculate_1q_p_value(ag_x1_, ag_x2_, Dinv_, Uinv__): + # anisotropic gaussian x1 and x2 + # calculate p-value with a 1-quadrant constraint + n_r = len(ag_x1_) + p1q_equ_r_ = np.zeros(n_r) + tmp_a = max(np.sqrt(Dinv_)) + tmp_b = min(np.sqrt(Dinv_)) + #print("Dinv = [", Dinv_[0], ", ", Dinv_[1], "]") + #print("first eig-vec = ", Uinv__[:,0]) + #print("second eig-vec = ", Uinv__[:,1]) + if Dinv_[0]<=Dinv_[1]: + Uinv__ = np.array([Uinv__[:,1], Uinv__[:,0]]).T + if (Uinv__[1,0] < 0) and (Uinv__[0,0]<0): + Uinv__[1,0]*=-1 + Uinv__[0,0]*=-1 + tmp_w = math.atan2(Uinv__[1,0], Uinv__[0,0]) + # print("Major axis = "+str(tmp_w*180/np.pi)+" deg \n\n") + tmp_gamma = math.atan2(1.0, 0.5*np.sin(2*tmp_w)*(tmp_b/tmp_a-tmp_a/tmp_b)) # angular formula only valid in 2-dimensions + for i in range(0, n_r): + ag_x1 = ag_x1_[i] + ag_x2 = ag_x2_[i] + p1q_equ = 1.0 + if (ag_x1>0) and (ag_x2>0): + tmp_x_0 = +np.cos(tmp_w)*ag_x1 + np.sin(tmp_w)*ag_x2 + tmp_x_1 = -np.sin(tmp_w)*ag_x1 + np.cos(tmp_w)*ag_x2 # R*S: [cos(w) -sin(w), sin(w) cos(w)]*[sx 0, 0 sy] rotation matrix * scaling matrix + tmp_y_0 = tmp_x_0/max(1e-12, tmp_a) + tmp_y_1 = tmp_x_1/max(1e-12, tmp_b) + tmp_y_r = np.sqrt(tmp_y_0**2+tmp_y_1**2) + p1q_equ = np.exp(-tmp_y_r**2/2)*tmp_gamma/(2*np.pi) + p1q_equ_r_[i] = p1q_equ + + return p1q_equ_r_,-np.log(p1q_equ_r_) + +def calculate_3q_p_value(ag_x1_, ag_x2_, Dinv_, Uinv__): + # anisotropic gaussian x1 and x2 + # calculate p-value with a 1-quadrant constraint + n_r = len(ag_x1_) + p1q_equ_r_ = np.zeros(n_r) + tmp_a = max(np.sqrt(Dinv_)) + tmp_b = min(np.sqrt(Dinv_)) + #print("Dinv = [", Dinv_[0], ", ", Dinv_[1], "]") + #print("first eig-vec = ", Uinv__[:,0]) + #print("second eig-vec = ", Uinv__[:,1]) + if Dinv_[0]<=Dinv_[1]: + Uinv__ = np.array([Uinv__[:,1], Uinv__[:,0]]).T + if (Uinv__[1,0] < 0) and (Uinv__[0,0]<0): + Uinv__[1,0]*=-1 + Uinv__[0,0]*=-1 + tmp_w = math.atan2(Uinv__[1,0], Uinv__[0,0]) + # print("Major axis = "+str(tmp_w*180/np.pi)+" deg \n\n") + tmp_gamma = math.atan2(1.0, 0.5*np.sin(2*tmp_w)*(tmp_b/tmp_a-tmp_a/tmp_b)) # angular formula only valid in 2-dimensions + for i in range(0, n_r): + ag_x1 = ag_x1_[i] + ag_x2 = ag_x2_[i] + p1q_equ = 1.0 + if (ag_x1>0) or (ag_x2>0): + tmp_x_0 = +np.cos(tmp_w)*ag_x1 + np.sin(tmp_w)*ag_x2 + tmp_x_1 = -np.sin(tmp_w)*ag_x1 + np.cos(tmp_w)*ag_x2 # R*S: [cos(w) -sin(w), sin(w) cos(w)]*[sx 0, 0 sy] rotation matrix * scaling matrix + tmp_y_0 = tmp_x_0/max(1e-12, tmp_a) + tmp_y_1 = tmp_x_1/max(1e-12, tmp_b) + tmp_y_r = np.sqrt(tmp_y_0**2+tmp_y_1**2) + p1q_equ = np.exp(-tmp_y_r**2/2)*tmp_gamma/(2*np.pi) + p1q_equ_r_[i] = p1q_equ + + return p1q_equ_r_,-np.log(p1q_equ_r_) + +def calculate_2dtm_pval(zscores, snrs, q=1): + """ + Calculate the 2D p-value for a given z-score and SNR using anisotropic Gaussian distribution. + Args: + zscore (numpy.ndarray): Array of z-scores. + snr (numpy.ndarray): Array of SNR values. + Returns: + numpy.ndarray: Array of (negative log) p-values. + """ + x1_ = zscores + x2_ = snrs + _, pro_x1_ = calculate_probit(x1_) + _, pro_x2_ = calculate_probit(x2_) + + Dinv_, Uinv__ = estimate_anisotropic_gaussian_for_probit(pro_x1_, pro_x2_) + if q==1: + pval,neg_log_p1q_equ_pro_ = calculate_1q_p_value(pro_x1_, pro_x2_, Dinv_, Uinv__) + if q==3: + pval,neg_log_p1q_equ_pro_ = calculate_3q_p_value(pro_x1_, pro_x2_, Dinv_, Uinv__) + return neg_log_p1q_equ_pro_ \ No newline at end of file