File size: 8,500 Bytes
6e4d123
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
"""Evaluate FASTA with a trained composition-features logistic regression (optional labeled CSV)."""
from __future__ import annotations

import argparse
import math
import os
import sys
from pathlib import Path

import joblib
import numpy as np
import pandas as pd

#
# =====================================================================
#  Default checkpoint: training output dir or repository `weights/composition_lr/`
#  (see evaluation/README.md). Must contain composition_lr_model.joblib (+ calibration CSVs).
#  Override via --model-dir.
# =====================================================================
_REPO_ROOT = Path(__file__).resolve().parents[1]
MODEL_CHECKPOINT_DIR = _REPO_ROOT / 'weights' / 'composition_lr'
#

_MODELS_DIR = _REPO_ROOT / 'models'
if str(_MODELS_DIR) not in sys.path:
    sys.path.insert(0, str(_MODELS_DIR))

from lr_sequence_composition_baseline import (  # noqa: E402
    FEATURE_NAMES,
    compute_composition_features,
    predict_proba,
    standard_amino_acids,
)

_SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
if _SCRIPT_DIR not in sys.path:
    sys.path.insert(0, _SCRIPT_DIR)

import eval_metrics as em  # noqa: E402


def resolve_model_dir(cli_value: str | None) -> str:
    if cli_value:
        return os.path.abspath(cli_value)
    return str(MODEL_CHECKPOINT_DIR.resolve())


def load_inference_operating_point(model_dir: str, model_name: str = 'composition_lr') -> tuple[float, float]:
    temperature = 1.0
    threshold = 0.5
    temp_path = os.path.join(model_dir, f'{model_name}_validation_temperature.csv')
    if os.path.isfile(temp_path):
        try:
            df = pd.read_csv(temp_path)
            if 'temperature' in df.columns and len(df) > 0:
                temperature = float(df['temperature'].iloc[0])
        except Exception as e:
            print(f'Warning: temperature file: {e}')
    summary_path = os.path.join(model_dir, f'{model_name}_validation_validation_diagnostics_summary.csv')
    if not os.path.isfile(summary_path):
        summary_path = os.path.join(model_dir, f'{model_name}_validation_diagnostics_summary.csv')
    if os.path.isfile(summary_path):
        try:
            sdf = pd.read_csv(summary_path)
            if 'optimal_threshold_f1' in sdf.columns and len(sdf) > 0:
                threshold = float(sdf['optimal_threshold_f1'].iloc[0])
            if 'temperature' in sdf.columns and len(sdf) > 0 and not os.path.isfile(temp_path):
                temperature = float(sdf['temperature'].iloc[0])
        except Exception as e:
            print(f'Warning: diagnostics summary: {e}')
    return temperature, threshold


def load_composition_model(model_dir: str):
    model_path = os.path.join(model_dir, 'composition_lr_model.joblib')
    if not os.path.isfile(model_path):
        raise FileNotFoundError(
            f'{model_path} not found. Re-run lr_sequence_composition_baseline.py to export '
            'composition_lr_model.joblib, or pass --model-dir to a directory that contains it.'
        )
    bundle = joblib.load(model_path)
    clf = bundle['clf']
    scaler = bundle['scaler']
    feature_names = list(bundle.get('feature_names', FEATURE_NAMES))
    if feature_names != FEATURE_NAMES:
        print(
            f'Warning: checkpoint feature_names differ from current script; '
            f'using checkpoint order ({len(feature_names)} features).'
        )
    return clf, scaler, feature_names


def parse_arguments():
    p = argparse.ArgumentParser(
        description='Evaluate FASTA with composition LR (gravy, pI, charge, complexity).',
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    p.add_argument('--fasta', required=True, help='Input FASTA file')
    p.add_argument(
        '--model-dir',
        default=None,
        help='Overrides MODEL_CHECKPOINT_DIR defined at the top of this script',
    )
    p.add_argument('--output', required=True, help='Output directory')
    p.add_argument(
        '--csv',
        default=None,
        help='Optional CSV with id column (variant/sequence_id/id) and label (highFRET/lowFRET)',
    )
    p.add_argument(
        '--no-detailed-metrics',
        action='store_true',
        help='With --csv: one-line summary only; skip diagnostic CSV/report output',
    )
    p.add_argument(
        '--include-features',
        action='store_true',
        help='Add per-sequence composition feature columns to the predictions CSV',
    )
    return p.parse_args()


def read_fasta_robust(fasta_path):
    h = None
    buf = []
    with open(fasta_path, 'r') as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            if line.startswith('>'):
                if h is not None:
                    s = ''.join(buf)
                    if s:
                        yield (h, s)
                h = line[1:].strip()
                buf = []
            elif h is not None:
                buf.append(line)
        if h is not None:
            s = ''.join(buf)
            if s:
                yield (h, s)


def main():
    args = parse_arguments()
    model_dir = resolve_model_dir(args.model_dir)
    detailed = args.csv and not args.no_detailed_metrics
    os.makedirs(args.output, exist_ok=True)

    print(f'Checkpoint directory: {model_dir}')
    clf, scaler, feature_names = load_composition_model(model_dir)
    temperature, threshold = load_inference_operating_point(model_dir)
    print(f'Operating point: temperature={temperature:.4f}, threshold={threshold:.4f}')

    sequence_ids: list[str] = []
    sequences: list[str] = []
    feature_rows: list[dict[str, float]] = []
    skipped: list[tuple[str, int, str]] = []

    for header, seq in read_fasta_robust(args.fasta):
        feats = compute_composition_features(seq)
        row = {name: feats[name] for name in feature_names}
        if any(math.isnan(v) for v in row.values()):
            skipped.append((header, len(standard_amino_acids(seq)), 'nan_composition_features'))
            continue
        sequence_ids.append(header)
        sequences.append(seq)
        feature_rows.append(row)

    for header, length, reason in skipped:
        print(f'Warning: skip {header} ({reason}, canonical_aa_len={length})')

    if not sequences:
        print('ERROR: No sequences to process.')
        return

    X = np.array([[row[n] for n in feature_names] for row in feature_rows], dtype=float)
    proba = predict_proba(clf, scaler, X)
    proba = em.apply_temperature_scaling(proba, temperature)

    results = pd.DataFrame(
        {
            'sequence_id': sequence_ids,
            'sequence_length': [len(standard_amino_acids(s)) for s in sequences],
            'prediction_probability': proba,
            'predicted_class': np.where(proba >= threshold, 'highFRET', 'lowFRET'),
        }
    )
    if args.include_features:
        feat_df = pd.DataFrame(feature_rows)
        results = pd.concat([results, feat_df], axis=1)

    out_predictions = os.path.join(args.output, 'composition_lr_predictions.csv')
    results.to_csv(out_predictions, index=False)
    print(f'Predictions saved: {out_predictions}')

    if skipped:
        pd.DataFrame(skipped, columns=['sequence_id', 'length', 'reason']).to_csv(
            os.path.join(args.output, 'skipped_sequences.csv'), index=False
        )

    if not args.csv or not os.path.isfile(args.csv):
        print('No labeled CSV; predictions only.')
        return

    try:
        _, _, label_map = em.build_label_map_int(pd.read_csv(args.csv))
    except ValueError as e:
        print(f'Skipping metrics: {e}')
        return

    y_true_list: list[int] = []
    y_prob_list: list[float] = []
    for sid in sequence_ids:
        lab = label_map.get(str(sid).strip())
        if lab is None:
            continue
        y_true_list.append(lab)
        y_prob_list.append(
            float(results.loc[results['sequence_id'] == sid, 'prediction_probability'].iloc[0])
        )
    if not y_true_list:
        print('No overlapping labels between FASTA IDs and CSV. Skipping metrics.')
        return

    y_true = np.asarray(y_true_list)
    y_prob = np.asarray(y_prob_list)
    y_pred = (y_prob >= threshold).astype(int)
    em.run_supervised_evaluation(
        y_true, y_prob, y_pred, 'composition_lr', args.output, threshold, detailed_metrics=detailed
    )
    if detailed:
        print(f'Detailed diagnostics written under {args.output}')


if __name__ == '__main__':
    main()