Coloring MSA based on associated value

plot_alignment.py

import numpy as np
import matplotlib.pyplot as plt
import biotite.sequence as seq
import biotite.sequence.align as align
import biotite.sequence.graphics as graphics

sequences = [
    seq.ProteinSequence(seq_str) for seq_str
    in ("BIQTITE", "TITANITE", "BISMITE", "IQLITE")
]
# Use random values for each sequence position for this example
np.random.seed(0)
values = [np.random.rand(len(sequence)) for sequence in sequences]

alignment, _, _, _ = align.align_multiple(
    sequences,
    matrix=align.SubstitutionMatrix.std_protein_matrix(),
    gap_penalty=-5,
    terminal_penalty=False
)


class ValuePlotter(graphics.LetterPlotter):

    def __init__(self, axes, values, cmap,
                 color_symbols=False, font_size=None, font_param=None):
        super().__init__(axes, color_symbols, font_size, font_param)
        self._values = values
        self._cmap = cmap
    
    def get_color(self, alignment, column_i, seq_i):
        seq_pos = alignment.trace[column_i, seq_i]
        if seq_pos == -1:
            # Gaps are white
            return (1, 1, 1)
        else:
            value = self._values[seq_i][seq_pos]
            return self._cmap(value)


fig, ax = plt.subplots(constrained_layout=True)
symbol_plotter = ValuePlotter(ax, values, plt.get_cmap("turbo"))
graphics.plot_alignment(
    ax, alignment, symbol_plotter, symbols_per_line=len(alignment)
)
plt.show()