140 lines · 4.9 KB · Python
RawDownload
#!/usr/bin/env python3"""draw_restriction_map.py - Annotate restriction enzyme cut sites on DNA sequences. Renders a linear or circular diagram showing restriction enzyme recognition sitesand cut positions using BioPython's Restriction module and matplotlib. Usage:    python draw_restriction_map.py --input sequence.gb --output restriction_map.png    python draw_restriction_map.py --input sequence.fasta --output map.png --enzymes EcoRI,BamHI,HindIII    python draw_restriction_map.py --input sequence.gb --output map.svg --linear --figwidth 16""" import argparseimport sys try:    from Bio import SeqIO    from Bio.Seq import Seq    from Bio.Restriction import RestrictionBatch, CommOnlyexcept ImportError:    print("Error: BioPython is required. Install with: pip install biopython", file=sys.stderr)    sys.exit(1) import matplotlibmatplotlib.use("Agg")import matplotlib.pyplot as pltimport matplotlib.patches as mpatchesimport numpy as np  # Colorblind-safe palette for enzyme labelsENZYME_COLORS = [    "#e74c3c", "#3498db", "#2ecc71", "#f39c12", "#9b59b6",    "#1abc9c", "#e67e22", "#34495e", "#d35400", "#8e44ad",    "#16a085", "#c0392b", "#2980b9", "#27ae60", "#f1c40f",]  def read_sequence(path):    ext = path.lower()    if ext.endswith((".gb", ".gbk", ".genbank")):        record = SeqIO.read(path, "genbank")    else:        record = SeqIO.read(path, "fasta")    return record  def find_cut_sites(seq, enzyme_names=None):    if enzyme_names:        batch = RestrictionBatch(enzyme_names)    else:        batch = CommOnly     results = batch.search(seq)     # Filter to enzymes that actually cut    sites = {}    for enzyme, positions in results.items():        if positions:            sites[str(enzyme)] = sorted(positions)    return sites  def main():    parser = argparse.ArgumentParser(description="Generate restriction enzyme site maps")    parser.add_argument("--input", required=True, help="GenBank or FASTA file path")    parser.add_argument("--output", required=True, help="Output image path (PNG/SVG/PDF)")    parser.add_argument("--enzymes", default=None,                        help="Comma-separated enzyme names (default: all common enzymes that cut)")    parser.add_argument("--max-enzymes", type=int, default=15,                        help="Max enzymes to display (default: 15)")    parser.add_argument("--linear", action="store_true", help="Force linear display")    parser.add_argument("--figwidth", type=float, default=14, help="Figure width in inches")    parser.add_argument("--dpi", type=int, default=300, help="Resolution in DPI")    parser.add_argument("--title", default=None, help="Figure title")    args = parser.parse_args()     record = read_sequence(args.input)    seq = record.seq    seq_len = len(seq)     enzyme_names = args.enzymes.split(",") if args.enzymes else None    sites = find_cut_sites(seq, enzyme_names)     if not sites:        print("No restriction enzyme cut sites found.", file=sys.stderr)        sys.exit(0)     # Sort by number of cuts (fewest first), limit display    sorted_enzymes = sorted(sites.keys(), key=lambda e: len(sites[e]))    if len(sorted_enzymes) > args.max_enzymes:        sorted_enzymes = sorted_enzymes[:args.max_enzymes]        print(f"Showing top {args.max_enzymes} enzymes (fewest cuts). "              f"Use --enzymes to specify.", file=sys.stderr)     # Plot    n_enzymes = len(sorted_enzymes)    fig_h = max(3, 1 + n_enzymes * 0.4)    fig, ax = plt.subplots(figsize=(args.figwidth, fig_h))     # Draw sequence backbone    ax.plot([0, seq_len], [0, 0], color="black", linewidth=2, zorder=1)     # Draw tick marks at each cut site    for i, enzyme in enumerate(sorted_enzymes):        positions = sites[enzyme]        color = ENZYME_COLORS[i % len(ENZYME_COLORS)]        y_offset = -(i + 1) * 0.35         for pos in positions:            ax.plot([pos, pos], [0, y_offset], color=color, linewidth=1, alpha=0.7)            ax.plot(pos, y_offset, "v", color=color, markersize=5)         # Label        ax.text(seq_len * 1.02, y_offset, f"{enzyme} ({len(positions)})",                fontsize=9, va="center", color=color, fontweight="bold")     # Axis formatting    ax.set_xlim(-seq_len * 0.02, seq_len * 1.15)    ax.set_ylim(-(n_enzymes + 1) * 0.35, 0.5)    ax.set_xlabel("Position (bp)", fontsize=11)    ax.set_yticks([])    ax.spines["top"].set_visible(False)    ax.spines["right"].set_visible(False)    ax.spines["left"].set_visible(False)     title = args.title or f"Restriction Map — {record.name} ({seq_len:,} bp)"    ax.set_title(title, fontsize=13, fontweight="bold", pad=10)     plt.tight_layout()    plt.savefig(args.output, dpi=args.dpi, bbox_inches="tight")    plt.close()     total_sites = sum(len(sites[e]) for e in sorted_enzymes)    print(f"Restriction map saved to {args.output} ({n_enzymes} enzymes, {total_sites} cut sites)")  if __name__ == "__main__":    main()