用 pandas 按样本 ID 连接蛋白质序列与标签，避免数据错位

python3 -m venv .venv
source .venv/bin/activate
python -m pip install "pandas==3.0.6"

python join_labels.py --help
python - <<'PY'
from pathlib import Path
Path("sequences.csv").write_text(
    "sample_id,sequence\n001,ACDE\n002,FGHI\nNA,KLMN\n",
    encoding="utf-8",
)
Path("labels.csv").write_text(
    "sample_id,label\nNA,positive\n001,negative\n002,positive\n",
    encoding="utf-8",
)
PY
python join_labels.py sequences.csv labels.csv joined.csv

OK: 3 sequences, 3 labels, 3 matched rows

sample_id,sequence,label
001,ACDE,negative
002,FGHI,positive
NA,KLMN,positive

"""Join tidy sequence and label CSVs by sample_id, with a strict 1:1 policy."""
import argparse
from pathlib import Path
import sys
import warnings

import pandas as pd


def read_table(path, value_column):
    with warnings.catch_warnings():
        warnings.simplefilter("error", pd.errors.ParserWarning)
        table = pd.read_csv(path, dtype="string", keep_default_na=False,
                            skip_blank_lines=False, index_col=False)
    columns = ["sample_id", value_column]
    if len(table.columns) != 2 or set(table.columns) != set(columns):
        raise ValueError(f"{path}: expected exactly {columns}")
    table = table[columns]
    if table.empty:
        raise ValueError(f"{path}: no data rows")
    for column in columns:
        values = table[column]
        if values.isna().any() or values.str.strip().eq("").any():
            raise ValueError(f"{path}: blank {column}")
    ids = table["sample_id"]
    if ids.ne(ids.str.strip()).any():
        raise ValueError(f"{path}: sample_id has surrounding whitespace")
    duplicates = ids[ids.duplicated(keep=False)].unique().tolist()
    if duplicates:
        raise ValueError(f"{path}: duplicate sample_id: {duplicates[:5]}")
    return table


def join_tables(sequences, labels):
    audit = sequences.merge(labels, on="sample_id", how="outer",
                            validate="one_to_one", indicator=True)
    missing = audit.loc[audit["_merge"].eq("left_only"), "sample_id"].tolist()
    extra = audit.loc[audit["_merge"].eq("right_only"), "sample_id"].tolist()
    if missing or extra:
        raise ValueError(f"ID mismatch: missing labels={missing[:5]}; "
                         f"labels without sequences={extra[:5]}")
    # Use the left join only after the full coverage audit has passed.
    return sequences.merge(labels, on="sample_id", how="left", sort=False,
                           validate="one_to_one")


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("sequences", type=Path)
    parser.add_argument("labels", type=Path)
    parser.add_argument("output", type=Path)
    args = parser.parse_args()
    try:
        sequences = read_table(args.sequences, "sequence")
        labels = read_table(args.labels, "label")
        joined = join_tables(sequences, labels)
        # Exclusive creation preserves earlier outputs and the input files.
        with args.output.open("x", encoding="utf-8", newline="") as handle:
            joined.to_csv(handle, index=False, lineterminator="\n")
        print(f"OK: {len(joined)} sequences, {len(labels)} labels, "
              f"{len(joined)} matched rows")
        return 0
    except (OSError, ValueError, pd.errors.ParserError,
            pd.errors.ParserWarning, pd.errors.EmptyDataError,
            pd.errors.MergeError) as error:
        print(f"ERROR: {error}", file=sys.stderr)
        return 2


if __name__ == "__main__":
    sys.exit(main())

python - <<'PY'
from pathlib import Path
Path("labels-missing.csv").write_text(
    "sample_id,label\nNA,positive\n001,negative\n",
    encoding="utf-8",
)
PY
python join_labels.py sequences.csv labels-missing.csv rejected.csv
printf 'exit code: %s\n' "$?"

ERROR: ID mismatch: missing labels=['002']; labels without sequences=[]
exit code: 2