Add explicit sequence-selection manifests

Helixer.py

@@ -17,6 +17,7 @@
 from helixer.prediction.HybridModel import HybridModel
 from helixer.export.exporter import HelixerFastaToH5Controller
 from helixer.core.helpers import get_log_dict
+from helixer.core.sequence_selection import parse_manifest, read_fasta, write_report, write_selected_fasta
 
 
 class HelixerParameterParser(ParameterParser):
@@ -27,6 +28,12 @@ def __init__(self, config_file_path: str = '') -> None:
         self.io_group.add_argument('--species', type=str, help='Species name.')
         self.io_group.add_argument('--temporary-dir', type=str,
                                    help='use supplied (instead of system default) for temporary directory')
+        self.io_group.add_argument('--sequence-selection', type=str,
+                                   help='TSV with exactly one true/false decision per FASTA seqid. Selection changes annotation scope.')
+        self.io_group.add_argument('--selection-report', type=str,
+                                   help='TSV path for exact full/selected workload and provenance accounting.')
+        self.io_group.add_argument('--selection-only', action='store_true',
+                                   help='Validate selection and write its report without exporter, model, or post-processing.')
 
         self.data_group.add_argument('--subsequence-length', type=int,
                                      help='How to slice the genomic sequence. Set moderately longer than length of '
@@ -106,6 +113,9 @@ def __init__(self, config_file_path: str = '') -> None:
             'edge_threshold': 0.1,
             'peak_threshold': 0.8,
             'min_coding_length': 60,
+            'sequence_selection': None,
+            'selection_report': None,
+            'selection_only': False,
         }
         self.defaults = {**self.defaults, **helixer_defaults}
 
@@ -122,6 +132,12 @@ def check_for_lineage_model(lineage: str, downloaded_model_path: str) -> str:
 
     def check_args(self, args: argparse.Namespace) -> None:
 
+        if args.selection_only:
+            assert args.sequence_selection is not None, '--selection-only requires --sequence-selection'
+            assert args.selection_report is not None, '--selection-only requires --selection-report'
+            assert args.subsequence_length is not None, '--selection-only requires --subsequence-length'
+            return
+
         if args.model_filepath is not None:
             print(f'overriding the lineage based model, '
                   f'with the manually specified {args.model_filepath}', file=sys.stderr)
@@ -191,6 +207,15 @@ def main() -> None:
     pp = HelixerParameterParser('config/helixer_config.yaml')
     args = pp.get_args()
     args.overlap = not args.no_overlap  # minor overlapping is a far better default for inference. Thus, this hack.
+    records = decisions = None
+    if args.sequence_selection:
+        records = read_fasta(args.fasta_path)
+        decisions = parse_manifest(args.sequence_selection, records)
+        if args.selection_report:
+            write_report(args.selection_report, records, decisions, args.subsequence_length, args.sequence_selection)
+        if args.selection_only:
+            logger.info('Selection-only validation complete; no exporter, model, or post-processing was launched.')
+            return
     # before we start, check if helixer_post_bin will (presumably) be able to run
     # first, is it there
     logger.info(colored('\nHelixer.py config:\n', 'yellow') + f'{pformat(vars(args))}\n')
@@ -238,7 +263,11 @@ def main() -> None:
         tmp_genome_h5_path = os.path.join(tmp_dirname, f'tmp_species_{args.species}.h5')
         tmp_pred_h5_path = os.path.join(tmp_dirname, f'tmp_predictions_{args.species}.h5')
 
-        controller = HelixerFastaToH5Controller(args.fasta_path, tmp_genome_h5_path)
+        fasta_path = args.fasta_path
+        if records is not None and decisions is not None:
+            fasta_path = os.path.join(tmp_dirname, 'selected_sequences.fa')
+            write_selected_fasta(fasta_path, records, decisions)
+        controller = HelixerFastaToH5Controller(fasta_path, tmp_genome_h5_path)
         # hard coded subsequence length due to how the models have been created
         controller.export_fasta_to_h5(chunk_size=args.subsequence_length, compression=args.compression,
                                       multiprocess=not args.no_multiprocess, species=args.species,

helixer/core/sequence_selection.py

@@ -0,0 +1,116 @@
+"""Strict, criterion-agnostic FASTA sequence-selection manifests."""
+
+from __future__ import annotations
+
+import csv
+import hashlib
+from collections import Counter
+from dataclasses import dataclass
+from pathlib import Path
+
+
+@dataclass(frozen=True)
+class FastaRecord:
+    seqid: str
+    sequence: str
+
+
+@dataclass(frozen=True)
+class SelectionDecision:
+    seqid: str
+    include: bool
+    reason: str
+
+
+def read_fasta(path: str) -> list[FastaRecord]:
+    records: list[FastaRecord] = []
+    seqid: str | None = None
+    sequence: list[str] = []
+    with open(path, encoding='utf-8') as handle:
+        for line_no, line in enumerate(handle, 1):
+            line = line.rstrip('\n\r')
+            if line.startswith('>'):
+                if seqid is not None:
+                    records.append(FastaRecord(seqid, ''.join(sequence)))
+                seqid = line[1:].split(None, 1)[0]
+                if not seqid:
+                    raise ValueError(f'empty FASTA identifier at line {line_no}')
+                sequence = []
+            elif seqid is None:
+                if line.strip():
+                    raise ValueError(f'FASTA sequence before first header at line {line_no}')
+            else:
+                sequence.append(line.strip())
+    if seqid is not None:
+        records.append(FastaRecord(seqid, ''.join(sequence)))
+    if not records:
+        raise ValueError(f'FASTA input {path!r} contains no records')
+    ids = [record.seqid for record in records]
+    duplicate = next((item for item, count in Counter(ids).items() if count > 1), None)
+    if duplicate:
+        raise ValueError(f"duplicate FASTA seqid {duplicate!r}")
+    return records
+
+
+def parse_manifest(path: str, records: list[FastaRecord]) -> list[SelectionDecision]:
+    manifest_bytes = Path(path).read_bytes()
+    try:
+        rows = list(csv.DictReader(manifest_bytes.decode('utf-8').splitlines(), delimiter='\t'))
+    except UnicodeDecodeError as error:
+        raise ValueError(f'selection manifest is not UTF-8: {error}') from error
+    if not rows or rows[0] is None:
+        raise ValueError('selection manifest has no decisions')
+    required = {'seqid', 'include'}
+    fieldnames = set(rows[0])
+    if not required.issubset(fieldnames) or not fieldnames.issubset({'seqid', 'include', 'reason'}):
+        raise ValueError('selection manifest columns must be seqid, include, and optional reason')
+    decisions: dict[str, SelectionDecision] = {}
+    fasta_ids = {record.seqid for record in records}
+    for row_no, row in enumerate(rows, 2):
+        seqid = (row.get('seqid') or '').strip()
+        value = (row.get('include') or '').strip().lower()
+        if not seqid:
+            raise ValueError(f'empty seqid in selection manifest row {row_no}')
+        if seqid in decisions:
+            raise ValueError(f'duplicate manifest seqid {seqid!r}')
+        if seqid not in fasta_ids:
+            raise ValueError(f'manifest seqid {seqid!r} is absent from FASTA')
+        if value not in {'true', 'false'}:
+            raise ValueError(f"invalid include value {value!r} for {seqid!r}; use true or false")
+        decisions[seqid] = SelectionDecision(seqid, value == 'true', (row.get('reason') or '').strip())
+    missing = fasta_ids - set(decisions)
+    if missing:
+        raise ValueError(f'manifest has no decision for FASTA seqid {sorted(missing)[0]!r}')
+    ordered = [decisions[record.seqid] for record in records]
+    if not any(item.include for item in ordered):
+        raise ValueError('selection manifest excludes every FASTA record')
+    return ordered
+
+
+def manifest_checksum(path: str) -> str:
+    return hashlib.sha256(Path(path).read_bytes()).hexdigest()
+
+
+def write_selected_fasta(path: str, records: list[FastaRecord], decisions: list[SelectionDecision]) -> None:
+    with open(path, 'w', encoding='utf-8') as handle:
+        for record, decision in zip(records, decisions):
+            if decision.include:
+                handle.write(f'>{record.seqid}\n{record.sequence}\n')
+
+
+def write_report(path: str, records: list[FastaRecord], decisions: list[SelectionDecision], chunk_size: int,
+                 manifest_path: str) -> None:
+    reason_counts = Counter(item.reason for item in decisions if item.reason)
+    with open(path, 'w', newline='', encoding='utf-8') as handle:
+        writer = csv.writer(handle, delimiter='\t')
+        writer.writerow(['metric', 'value'])
+        for label, subset in [('full', list(zip(records, decisions))), ('selected', [(r, d) for r, d in zip(records, decisions) if d.include])]:
+            bp = sum(len(record.sequence) for record, _ in subset)
+            rows = sum(2 * ((len(record.sequence) + chunk_size - 1) // chunk_size) for record, _ in subset)
+            writer.writerow([f'{label}_records', len(subset)])
+            writer.writerow([f'{label}_bp', bp])
+            writer.writerow([f'{label}_two_strand_fixed_window_rows', rows])
+            writer.writerow([f'{label}_padded_positions', rows * chunk_size])
+        writer.writerow(['manifest_sha256', manifest_checksum(manifest_path)])
+        for reason, count in sorted(reason_counts.items()):
+            writer.writerow([f'reason_count:{reason}', count])

helixer/tests/test_sequence_selection.py

@@ -0,0 +1,38 @@
+import pytest
+
+from helixer.core.sequence_selection import parse_manifest, read_fasta, write_report, write_selected_fasta
+
+
+def write(path, text):
+    path.write_text(text, encoding='utf-8')
+    return str(path)
+
+
+def test_selected_fasta_preserves_order_identifiers_and_sequences(tmp_path):
+    fasta = write(tmp_path / 'input.fa', '>first description\nACGT\n>second\nTT\n')
+    manifest = write(tmp_path / 'selection.tsv', 'seqid\tinclude\treason\nfirst\tfalse\tshort\nsecond\ttrue\tkeep\n')
+    records = read_fasta(fasta)
+    decisions = parse_manifest(manifest, records)
+    selected = tmp_path / 'selected.fa'
+    write_selected_fasta(selected, records, decisions)
+    assert selected.read_text() == '>second\nTT\n'
+    report = tmp_path / 'report.tsv'
+    write_report(report, records, decisions, 4, manifest)
+    text = report.read_text()
+    assert 'full_two_strand_fixed_window_rows\t4' in text
+    assert 'selected_two_strand_fixed_window_rows\t2' in text
+    assert 'selected_padded_positions\t8' in text
+
+
+@pytest.mark.parametrize('manifest, message', [
+    ('seqid\tinclude\na\ttrue\na\tfalse\nb\ttrue\n', 'duplicate'),
+    ('seqid\tinclude\na\ttrue\n', 'no decision'),
+    ('seqid\tinclude\na\tyes\nb\ttrue\n', 'invalid include'),
+    ('seqid\tinclude\na\tfalse\nb\tfalse\n', 'excludes every'),
+    ('seqid\tinclude\na\ttrue\nunknown\tfalse\nb\ttrue\n', 'absent'),
+])
+def test_manifest_rejects_unsafe_decisions(tmp_path, manifest, message):
+    records = read_fasta(write(tmp_path / 'input.fa', '>a\nA\n>b\nT\n'))
+    path = write(tmp_path / 'selection.tsv', manifest)
+    with pytest.raises(ValueError, match=message):
+        parse_manifest(path, records)