Buckets:
| #!/usr/bin/env python3 | |
| """ | |
| seq_fragmentizer.py | |
| Fragments sequences respecting a pre-made train/test split, preventing data leakage. | |
| Usage: | |
| python3 seq_fragmentizer.py \\ | |
| --split erv_seq_split.csv \\ | |
| --input hg38_ERV_100bp.fa \\ | |
| --min-length 100 \\ | |
| --frag-min 150 --frag-max 350 \\ | |
| --train-count 1000 --test-count 1000 \\ | |
| --seed 42 \\ | |
| --out-prefix ERV | |
| Inputs: | |
| --split CSV file with columns: sequence_name, sequence_length, group (train/test) | |
| --input FASTA file with sequences to fragment | |
| Outputs: | |
| <out-prefix>_train.fa -- training FASTA | |
| <out-prefix>_test.fa -- testing FASTA | |
| <out-prefix>_qc_report.txt -- QC statistics | |
| Quality Control (applied BEFORE fragmentation): | |
| 1. Minimum length filter (--min-length, default 100bp) | |
| 2. Ambiguous nucleotide filter (remove sequences with 'N') | |
| 3. Invalid base filter (keep only A/C/G/T/a/c/g/t) | |
| Fragmentation: | |
| - Respects train/test partition from split file (NO data leakage) | |
| - Randomly samples windows of length [frag-min, frag-max] | |
| - Same sequence can be sampled multiple times with different positions | |
| - Each fragment guaranteed N-free and valid-base-only | |
| """ | |
| import argparse | |
| import random | |
| import sys | |
| def read_fasta(path): | |
| """Yield (header, sequence) tuples. Sequence is uppercase, newline-free.""" | |
| header = None | |
| seq_chunks = [] | |
| with open(path) as fh: | |
| for line in fh: | |
| line = line.rstrip("\n") | |
| if not line: | |
| continue | |
| if line.startswith(">"): | |
| if header is not None: | |
| yield header, "".join(seq_chunks).upper() | |
| header = line[1:].strip() | |
| seq_chunks = [] | |
| else: | |
| seq_chunks.append(line) | |
| if header is not None: | |
| yield header, "".join(seq_chunks).upper() | |
| def write_fasta(records, path, wrap=70): | |
| """Write (header, seq) tuples to FASTA.""" | |
| with open(path, "w") as out: | |
| for header, seq in records: | |
| out.write(f">{header}\n") | |
| for i in range(0, len(seq), wrap): | |
| out.write(seq[i:i + wrap] + "\n") | |
| def read_split(path): | |
| """ | |
| Read train/test split CSV. Expected format: | |
| sequence_name,sequence_length,group | |
| chr1:100-200,100,train | |
| chr1:300-400,100,test | |
| ... | |
| Returns dict: {sequence_name -> 'train' or 'test'} | |
| """ | |
| split_map = {} | |
| with open(path) as fh: | |
| header = next(fh).strip() | |
| for line in fh: | |
| line = line.strip() | |
| if not line: | |
| continue | |
| parts = line.split(",") | |
| if len(parts) < 3: | |
| raise ValueError(f"Invalid split line (expected 3 columns): {line}") | |
| seq_id = parts[0].strip() | |
| partition = parts[2].strip().lower() | |
| if partition not in ("train", "test"): | |
| raise ValueError(f"Invalid partition '{partition}' for {seq_id} (must be 'train' or 'test')") | |
| split_map[seq_id] = partition | |
| return split_map | |
| def get_sequence_id(header): | |
| """Extract sequence ID from FASTA header (first token before space/pipe).""" | |
| return header.split()[0].split("|")[0] | |
| def has_valid_bases_only(seq): | |
| """Check if sequence contains only A/C/G/T (case-insensitive).""" | |
| valid = set("ACGT") | |
| return all(base in valid for base in seq) | |
| def has_n(seq): | |
| """Check if sequence contains 'N' (ambiguous nucleotide).""" | |
| return "N" in seq | |
| def random_window(seq, length, rng, max_tries=20): | |
| """ | |
| Return a random N-free, valid-base-only substring of `length` from seq, or None. | |
| Tries up to max_tries random positions before giving up. | |
| """ | |
| if len(seq) < length: | |
| return None | |
| for _ in range(max_tries): | |
| start = rng.randint(0, len(seq) - length) | |
| frag = seq[start:start + length] | |
| if not has_n(frag) and has_valid_bases_only(frag): | |
| return frag, start | |
| return None | |
| def qc_filter_sequences(records, min_length): | |
| """ | |
| Apply QC filters BEFORE fragmentation: | |
| 1. Minimum length (>= min_length) | |
| 2. No 'N' (ambiguous nucleotides) | |
| 3. Valid bases only (A/C/G/T) | |
| Returns (filtered_records, qc_stats_dict). | |
| """ | |
| stats = { | |
| "input": 0, | |
| "too_short": 0, | |
| "has_n": 0, | |
| "invalid_bases": 0, | |
| "passed": 0, | |
| } | |
| filtered = [] | |
| for header, seq in records: | |
| stats["input"] += 1 | |
| L = len(seq) | |
| if L < min_length: | |
| stats["too_short"] += 1 | |
| continue | |
| if has_n(seq): | |
| stats["has_n"] += 1 | |
| continue | |
| if not has_valid_bases_only(seq): | |
| stats["invalid_bases"] += 1 | |
| continue | |
| filtered.append((header, seq)) | |
| stats["passed"] += 1 | |
| return filtered, stats | |
| def fragment_partition(partition_records, target_count, frag_min, frag_max, rng): | |
| """ | |
| Create fragments for a single partition (train or test). | |
| Randomly sample windows of length [frag_min, frag_max] from partition_records. | |
| Continues sampling (with replacement) until target_count is reached or gives up. | |
| """ | |
| if not partition_records: | |
| return [], 0 | |
| frags = [] | |
| attempts = 0 | |
| max_attempts = target_count * 50 # generous retry budget | |
| while len(frags) < target_count and attempts < max_attempts: | |
| attempts += 1 | |
| header, seq = rng.choice(partition_records) | |
| frag_len = rng.randint(frag_min, frag_max) | |
| result = random_window(seq, frag_len, rng) | |
| if result: | |
| frag, start = result | |
| frag_id = f"{get_sequence_id(header)}|{start}-{start+frag_len}|{frag_len}bp" | |
| frags.append((frag_id, frag)) | |
| n_short = max(0, target_count - len(frags)) | |
| return frags, n_short | |
| def main(): | |
| ap = argparse.ArgumentParser( | |
| description=__doc__, | |
| formatter_class=argparse.RawDescriptionHelpFormatter | |
| ) | |
| ap.add_argument("--split", required=True, | |
| help="CSV file with train/test partition (columns: sequence_name,sequence_length,group)") | |
| ap.add_argument("--input", required=True, | |
| help="Input FASTA file to fragment") | |
| ap.add_argument("--min-length", type=int, default=100, | |
| help="Discard input sequences shorter than this (default 100bp)") | |
| ap.add_argument("--frag-min", type=int, required=True, | |
| help="Minimum fragment length (e.g., 150)") | |
| ap.add_argument("--frag-max", type=int, required=True, | |
| help="Maximum fragment length (e.g., 350)") | |
| ap.add_argument("--train-count", type=int, required=True, | |
| help="Target number of fragments for training set") | |
| ap.add_argument("--test-count", type=int, required=True, | |
| help="Target number of fragments for test set") | |
| ap.add_argument("--seed", type=int, default=42, | |
| help="Random seed for reproducibility (default 42)") | |
| ap.add_argument("--out-prefix", required=True, | |
| help="Prefix for output files") | |
| args = ap.parse_args() | |
| if args.frag_min > args.frag_max: | |
| print("ERROR: --frag-min must be <= --frag-max", file=sys.stderr) | |
| sys.exit(1) | |
| rng = random.Random(args.seed) | |
| # Read split file | |
| print(f"Reading split file: {args.split}", file=sys.stderr) | |
| split_map = read_split(args.split) | |
| print(f" Found {len(split_map)} sequences in split", file=sys.stderr) | |
| # Read and parse FASTA | |
| print(f"Reading input FASTA: {args.input}", file=sys.stderr) | |
| all_records = list(read_fasta(args.input)) | |
| print(f" Loaded {len(all_records)} sequences", file=sys.stderr) | |
| # Map sequences to partition | |
| train_records = [] | |
| test_records = [] | |
| missing_in_split = [] | |
| for header, seq in all_records: | |
| seq_id = get_sequence_id(header) | |
| if seq_id not in split_map: | |
| missing_in_split.append(seq_id) | |
| continue | |
| partition = split_map[seq_id] | |
| if partition == "train": | |
| train_records.append((header, seq)) | |
| elif partition == "test": | |
| test_records.append((header, seq)) | |
| if missing_in_split: | |
| print(f"WARNING: {len(missing_in_split)} sequences not found in split file " | |
| f"(first 10: {missing_in_split[:10]})", file=sys.stderr) | |
| print(f" Before QC: {len(train_records)} train, {len(test_records)} test", file=sys.stderr) | |
| # QC filter (before fragmentation) | |
| print(f"Applying QC filters (min-length={args.min_length}bp)...", file=sys.stderr) | |
| train_qc, train_stats = qc_filter_sequences(train_records, args.min_length) | |
| test_qc, test_stats = qc_filter_sequences(test_records, args.min_length) | |
| print(f" Train QC: {train_stats['passed']}/{train_stats['input']} passed " | |
| f"(too_short={train_stats['too_short']}, has_N={train_stats['has_n']}, " | |
| f"invalid_bases={train_stats['invalid_bases']})", file=sys.stderr) | |
| print(f" Test QC: {test_stats['passed']}/{test_stats['input']} passed " | |
| f"(too_short={test_stats['too_short']}, has_N={test_stats['has_n']}, " | |
| f"invalid_bases={test_stats['invalid_bases']})", file=sys.stderr) | |
| # Fragment each partition independently | |
| print(f"Fragmenting (range {args.frag_min}-{args.frag_max}bp)...", file=sys.stderr) | |
| train_frags, train_short = fragment_partition( | |
| train_qc, args.train_count, args.frag_min, args.frag_max, rng | |
| ) | |
| test_frags, test_short = fragment_partition( | |
| test_qc, args.test_count, args.frag_min, args.frag_max, rng | |
| ) | |
| print(f" Train: {len(train_frags)}/{args.train_count} fragments " | |
| f"({train_short} short of target)", file=sys.stderr) | |
| print(f" Test: {len(test_frags)}/{args.test_count} fragments " | |
| f"({test_short} short of target)", file=sys.stderr) | |
| # Write output FASTA files | |
| train_path = f"{args.out_prefix}_train.fa" | |
| test_path = f"{args.out_prefix}_test.fa" | |
| write_fasta(train_frags, train_path) | |
| write_fasta(test_frags, test_path) | |
| print(f" Wrote {train_path}", file=sys.stderr) | |
| print(f" Wrote {test_path}", file=sys.stderr) | |
| # Write QC report | |
| report_path = f"{args.out_prefix}_qc_report.txt" | |
| with open(report_path, "w") as out: | |
| out.write("=" * 70 + "\n") | |
| out.write("QC REPORT\n") | |
| out.write("=" * 70 + "\n\n") | |
| out.write("INPUT\n") | |
| out.write(f" Split file: {args.split}\n") | |
| out.write(f" Input FASTA: {args.input}\n") | |
| out.write(f" Sequences in split: {len(split_map)}\n") | |
| out.write(f" Sequences loaded: {len(all_records)}\n") | |
| out.write(f" Sequences missing from split: {len(missing_in_split)}\n\n") | |
| out.write("BEFORE QC FILTERING\n") | |
| out.write(f" Train sequences: {len(train_records)}\n") | |
| out.write(f" Test sequences: {len(test_records)}\n\n") | |
| out.write("TRAIN QC FILTERING (--min-length={})bp\n".format(args.min_length)) | |
| out.write(f" Input: {train_stats['input']}\n") | |
| out.write(f" Filtered - too short: {train_stats['too_short']}\n") | |
| out.write(f" Filtered - contains N: {train_stats['has_n']}\n") | |
| out.write(f" Filtered - invalid bases: {train_stats['invalid_bases']}\n") | |
| out.write(f" Passed: {train_stats['passed']}\n\n") | |
| out.write("TEST QC FILTERING (--min-length={}bp)\n".format(args.min_length)) | |
| out.write(f" Input: {test_stats['input']}\n") | |
| out.write(f" Filtered - too short: {test_stats['too_short']}\n") | |
| out.write(f" Filtered - contains N: {test_stats['has_n']}\n") | |
| out.write(f" Filtered - invalid bases: {test_stats['invalid_bases']}\n") | |
| out.write(f" Passed: {test_stats['passed']}\n\n") | |
| out.write("FRAGMENTATION ({}--{}bp)\n".format(args.frag_min, args.frag_max)) | |
| out.write(f" Train - target: {args.train_count}, generated: {len(train_frags)}\n") | |
| out.write(f" Test - target: {args.test_count}, generated: {len(test_frags)}\n") | |
| out.write(f" Random seed: {args.seed}\n\n") | |
| out.write("OUTPUT\n") | |
| out.write(f" Train FASTA: {train_path}\n") | |
| out.write(f" Test FASTA: {test_path}\n") | |
| out.write(f" QC Report: {report_path}\n") | |
| print(f" Wrote {report_path}", file=sys.stderr) | |
| print("Done.", file=sys.stderr) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 12.4 kB
- Xet hash:
- 7b90436ab9c492f5aa9a42ba128aceaf45633b5d3b42657dfcea4809eb25103b
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.