"""
Convert HapMap (.hmp.txt) to PLINK binary (.bed/.bim/.fam)
Handles IUPAC ambiguity codes by treating them as missing.
"""
import sys
import os
from array import array

def hmp_to_plink(hmp_path, out_prefix):
    print(f"Reading {hmp_path}...")

    with open(hmp_path, 'r') as f:
        header = f.readline().strip().split('\t')
        samples = header[11:]
        n_samples = len(samples)

    print(f"  {n_samples} samples")

    # IUPAC -> allele mapping
    iupac_map = {
        'A': 'A', 'C': 'C', 'G': 'G', 'T': 'T',
        'R': '0', 'Y': '0', 'S': '0', 'W': '0',
        'K': '0', 'M': '0', 'N': '0', '-': '0'
    }

    # --- Write .tfam ---
    with open(f"{out_prefix}.tfam", 'w') as f:
        for i, sid in enumerate(samples):
            f.write(f"{sid}\t{sid}\t0\t0\t0\t-9\n")

    # --- Write .tped ---
    # TPED format: chrom rsid genetic_dist pos <genotypes...>
    # Each SNP = 1 row, genotypes = 2 columns per sample
    n_snps = 0
    with open(hmp_path, 'r') as fin, open(f"{out_prefix}.tped", 'w') as fout:
        fin.readline()  # skip header
        for line in fin:
            cols = line.strip().split('\t')
            rs = cols[0]
            chrom = cols[2]
            pos = cols[3]

            fout.write(f"{chrom}\t{rs}\t0\t{pos}")

            for j in range(11, len(cols)):
                gt = cols[j]
                allele = iupac_map.get(gt, '0')
                fout.write(f"\t{allele}\t{allele}")

            fout.write("\n")
            n_snps += 1

    print(f"  {n_snps} SNPs written to {out_prefix}.tped, {out_prefix}.tfam")

    # --- Run PLINK to convert to binary ---
    import subprocess
    plink = r"D:\tools\Plink\plink.exe"

    cmd = [
        plink,
        "--tfile", out_prefix,
        "--make-bed",
        "--out", out_prefix,
        "--chr-set", "10",
        "--allow-extra-chr"
    ]
    print(f"Running: {' '.join(cmd)}")
    result = subprocess.run(cmd, capture_output=True, text=True)
    print(result.stdout)
    if result.returncode != 0:
        print(result.stderr)
        sys.exit(1)

    # --- Verify ---
    for ext in ['.bed', '.bim', '.fam']:
        fpath = f"{out_prefix}{ext}"
        if os.path.exists(fpath):
            size_mb = os.path.getsize(fpath) / (1024*1024)
            print(f"  {fpath}: {size_mb:.1f} MB")
        else:
            print(f"  ERROR: {fpath} not created!")

    print("\nDone!")


if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("Usage: python hmp2plink.py <input.hmp.txt> <output_prefix>")
        sys.exit(1)

    hmp_to_plink(sys.argv[1], sys.argv[2])

使用方法

python hmp2plink.py <输入.hmp.txt> <输出目录>
python C:\Users\zhang\.easyclaw\workspace\hmp2plink.py D:\data\houzhixin\CML\190821.hmp.txt D:\data\houzhixin\CML\plink\190821