270 lines
8.9 KiB
Python
270 lines
8.9 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import argparse
|
|
import csv
|
|
from collections import Counter
|
|
from pathlib import Path
|
|
|
|
|
|
LOCUS_HEADER = [
|
|
"OG",
|
|
"sequence_count",
|
|
"independent_event_count",
|
|
"max_single_sequence_unique_events",
|
|
"carrier_pattern",
|
|
"severe",
|
|
]
|
|
|
|
|
|
def read_fasta(fasta_path: Path) -> dict[str, str]:
|
|
sequences: dict[str, list[str]] = {}
|
|
sequence_id: str | None = None
|
|
|
|
with fasta_path.open() as fasta:
|
|
for line_number, raw_line in enumerate(fasta, start=1):
|
|
line = raw_line.strip()
|
|
if not line:
|
|
continue
|
|
if line.startswith(">"):
|
|
sequence_id = line[1:].split()[0]
|
|
if not sequence_id:
|
|
raise ValueError(
|
|
f"{fasta_path}:{line_number}: empty FASTA identifier"
|
|
)
|
|
if sequence_id in sequences:
|
|
raise ValueError(
|
|
f"{fasta_path}:{line_number}: duplicate FASTA identifier "
|
|
f"{sequence_id!r}"
|
|
)
|
|
sequences[sequence_id] = []
|
|
elif sequence_id is None:
|
|
raise ValueError(
|
|
f"{fasta_path}:{line_number}: sequence before first FASTA header"
|
|
)
|
|
else:
|
|
sequences[sequence_id].append(line)
|
|
|
|
if not sequences:
|
|
raise ValueError(f"{fasta_path}: no FASTA sequences found")
|
|
|
|
joined = {name: "".join(parts) for name, parts in sequences.items()}
|
|
lengths = {len(sequence) for sequence in joined.values()}
|
|
if len(lengths) != 1:
|
|
raise ValueError(f"{fasta_path}: aligned sequences have different lengths")
|
|
return joined
|
|
|
|
|
|
def event_mode(carrier_count: int, sequence_count: int) -> str:
|
|
if carrier_count == 1:
|
|
return "single_species"
|
|
if carrier_count == sequence_count - 1:
|
|
return "all_but_one"
|
|
if carrier_count == sequence_count:
|
|
return "all_taxa"
|
|
return "branch_shared"
|
|
|
|
|
|
def locus_pattern(
|
|
event_modes: list[str],
|
|
event_count: int,
|
|
max_unique_events: int,
|
|
) -> str:
|
|
modes = set(event_modes)
|
|
if max_unique_events >= 2:
|
|
return "single_species_scattered_multiple_events"
|
|
if event_count == 1 and modes == {"single_species"}:
|
|
return "single_species_single_event"
|
|
if modes == {"all_but_one"}:
|
|
return "all_but_one"
|
|
if modes == {"branch_shared"}:
|
|
return "branch_shared"
|
|
if modes == {"all_taxa"}:
|
|
return "all_taxa"
|
|
if not modes:
|
|
return "no_frameshift_detected"
|
|
return "mixed"
|
|
|
|
|
|
def merge_overlapping_triplets(
|
|
triplets: list[tuple[tuple[int, int, int], str]],
|
|
) -> list[tuple[tuple[int, ...], list[str]]]:
|
|
components: list[tuple[set[int], set[str]]] = []
|
|
|
|
for columns, carrier in triplets:
|
|
column_set = set(columns)
|
|
matching = [
|
|
index
|
|
for index, (component_columns, _) in enumerate(components)
|
|
if component_columns & column_set
|
|
]
|
|
|
|
if not matching:
|
|
components.append((column_set, {carrier}))
|
|
continue
|
|
|
|
merged_columns = set(column_set)
|
|
merged_carriers = {carrier}
|
|
for index in reversed(matching):
|
|
component_columns, component_carriers = components.pop(index)
|
|
merged_columns.update(component_columns)
|
|
merged_carriers.update(component_carriers)
|
|
components.append((merged_columns, merged_carriers))
|
|
|
|
return [
|
|
(tuple(sorted(columns)), sorted(carriers))
|
|
for columns, carriers in components
|
|
]
|
|
|
|
|
|
def classify_fasta(fasta_path: Path) -> tuple[list[str], list[list[str]]]:
|
|
sequences = read_fasta(fasta_path)
|
|
sequence_ids = list(sequences)
|
|
sequence_count = len(sequence_ids)
|
|
unique_events_by_sequence: Counter[str] = Counter()
|
|
event_rows: list[list[str]] = []
|
|
modes: list[str] = []
|
|
complete_triplets: list[tuple[tuple[int, int, int], str]] = []
|
|
events: list[tuple[tuple[int, ...], list[str]]] = []
|
|
|
|
for sequence_id, sequence in sequences.items():
|
|
marker_columns = [
|
|
column_index + 1
|
|
for column_index, character in enumerate(sequence)
|
|
if character == "!"
|
|
]
|
|
for index in range(0, len(marker_columns), 3):
|
|
event_columns = tuple(marker_columns[index : index + 3])
|
|
if len(event_columns) == 3:
|
|
complete_triplets.append((event_columns, sequence_id))
|
|
else:
|
|
# A trailing 1-2 markers still form one unique event, but cannot
|
|
# be classified as shared because they are not a complete triplet.
|
|
events.append((event_columns, [sequence_id]))
|
|
|
|
events.extend(merge_overlapping_triplets(complete_triplets))
|
|
events.sort(key=lambda event: (event[0], event[1]))
|
|
|
|
for event_columns, carriers in events:
|
|
carriers = [
|
|
sequence_id for sequence_id in sequence_ids if sequence_id in carriers
|
|
]
|
|
mode = event_mode(len(carriers), sequence_count)
|
|
modes.append(mode)
|
|
if len(carriers) == 1:
|
|
unique_events_by_sequence[carriers[0]] += 1
|
|
event_rows.append(
|
|
[
|
|
fasta_path.stem,
|
|
",".join(str(column) for column in event_columns),
|
|
str(len(carriers)),
|
|
",".join(carriers),
|
|
mode,
|
|
]
|
|
)
|
|
|
|
event_count = len(event_rows)
|
|
max_unique_events = max(unique_events_by_sequence.values(), default=0)
|
|
severe = event_count >= 3 or max_unique_events >= 2
|
|
locus_row = [
|
|
fasta_path.stem,
|
|
str(sequence_count),
|
|
str(event_count),
|
|
str(max_unique_events),
|
|
locus_pattern(modes, event_count, max_unique_events),
|
|
"yes" if severe else "no",
|
|
]
|
|
return locus_row, event_rows
|
|
|
|
|
|
def fasta_paths_from_list(list_path: Path) -> list[Path]:
|
|
paths: list[Path] = []
|
|
seen: set[Path] = set()
|
|
with list_path.open() as fasta_list:
|
|
for line_number, raw_line in enumerate(fasta_list, start=1):
|
|
value = raw_line.strip()
|
|
if not value:
|
|
continue
|
|
path = Path(value)
|
|
if not path.is_absolute():
|
|
path = list_path.parent / path
|
|
path = path.resolve()
|
|
if not path.is_file():
|
|
raise FileNotFoundError(
|
|
f"{list_path}:{line_number}: FASTA file not found: {path}"
|
|
)
|
|
if path not in seen:
|
|
paths.append(path)
|
|
seen.add(path)
|
|
return paths
|
|
|
|
|
|
def write_results(
|
|
output_dir: Path,
|
|
locus_rows: list[list[str]],
|
|
event_rows: list[list[str]],
|
|
) -> None:
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
with (output_dir / "frameshift_events.tsv").open("w", newline="") as output:
|
|
writer = csv.writer(output, delimiter="\t", lineterminator="\n")
|
|
writer.writerow(["OG", "alignment_column", "carrier_count", "carriers", "mode"])
|
|
writer.writerows(event_rows)
|
|
|
|
with (output_dir / "frameshift_locus_classification.tsv").open(
|
|
"w", newline=""
|
|
) as output:
|
|
writer = csv.writer(output, delimiter="\t", lineterminator="\n")
|
|
writer.writerow(LOCUS_HEADER)
|
|
writer.writerows(locus_rows)
|
|
|
|
severe_rows = [row for row in locus_rows if row[-1] == "yes"]
|
|
with (output_dir / "frameshift_severe_subset.tsv").open(
|
|
"w", newline=""
|
|
) as output:
|
|
writer = csv.writer(output, delimiter="\t", lineterminator="\n")
|
|
writer.writerow(LOCUS_HEADER)
|
|
writer.writerows(severe_rows)
|
|
|
|
pattern_counts = Counter(row[4] for row in locus_rows)
|
|
with (output_dir / "frameshift_classification_summary.tsv").open(
|
|
"w", newline=""
|
|
) as output:
|
|
writer = csv.writer(output, delimiter="\t", lineterminator="\n")
|
|
writer.writerow(["metric", "count"])
|
|
writer.writerow(["total_FS_loci", len(locus_rows)])
|
|
writer.writerow(["severe_loci", len(severe_rows)])
|
|
for pattern, count in sorted(pattern_counts.items()):
|
|
writer.writerow([pattern, count])
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(
|
|
description="Classify frameshift loci by event count and carrier pattern."
|
|
)
|
|
parser.add_argument(
|
|
"--fasta-list",
|
|
required=True,
|
|
type=Path,
|
|
help="File containing one aligned FASTA path per line.",
|
|
)
|
|
parser.add_argument(
|
|
"--output-dir",
|
|
required=True,
|
|
type=Path,
|
|
help="Directory for classification tables.",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
fasta_paths = fasta_paths_from_list(args.fasta_list.resolve())
|
|
locus_rows: list[list[str]] = []
|
|
event_rows: list[list[str]] = []
|
|
for fasta_path in fasta_paths:
|
|
locus_row, rows = classify_fasta(fasta_path)
|
|
locus_rows.append(locus_row)
|
|
event_rows.extend(rows)
|
|
write_results(args.output_dir, locus_rows, event_rows)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|