Add Contact Map with Triangle Attention Overlay

[salammurillo]

Addresses Issue #7

Adds distance map, contact map, and triangle attention over contact map visualization to show distances and triangle attention weights between residue pairs. The purpose of this visualization is to investigate the model's reasoning and attention weights across all heads and specified layers, and compare it to the physical structure of the protein.

Files Added

contact_map.py runs an inference using the input fasta and alignment directories on the protein and captures the triangle attention weights for the specified layers for all heads. then builds a contact map and overlays the attention weights per residue pair
generate_contact_map.py CLI utility to run contact_map.py and validates all inputs

How to Use

Run generate_contact_map.py using the following arguments'

Required:

• fasta file path
• alignment directory

Optional:

• region (Residues i - j)
• contact threshold (default 0.8)
• pdb_path
• layers (defaults to 38-47)

For example:

python generate_contact_map.py \
        examples/monomer/fasta_dir/6kwc.fasta \
        examples/monomer/alignments/6KWC_1 \
        --pdb_path outputs/predictions/6KWC_1_model_1_ptm_relaxed.pdb \
        --region 30 90

Produced Image: