Simple python script to generate an image of amino acid sequence

amino_acid_plot.py

# make your imports
import matplotlib.pyplot as plt
import matplotlib
import numpy as np


# Based on  this image
# https://www.google.com/imgres?imgurl=https%3A%2F%2Fqph.fs.quoracdn.net%2Fmain-qimg-88f77484a87d55f1cd615eb2ca2a8441&imgrefurl=https%3A%2F%2Fwww.quora.com%2FWhat-makes-some-amino-acids-polar-and-nonpolar&tbnid=I9Lczh2TViC1rM&vet=12ahUKEwi1qfv8x_fzAhUIr6QKHfN-A0cQMygDegUIARC4AQ..i&docid=Kf-tNnUawtUyBM&w=624&h=769&q=polar%20and%20non%20polar%20amino%20acids&client=firefox-b-d&ved=2ahUKEwi1qfv8x_fzAhUIr6QKHfN-A0cQMygDegUIARC4AQ
NEGATIVE_AA = "DE"
HYDROPHILIC_AA = "AVILMFYW"
POSITIVE_AA = "RHK"
POLAR_AA = "STNQ"
OTHER_AA = "CUGP"

# Create mapping between amino acid and index
AA_MAP = {}
for index, sequence in ((-1, NEGATIVE_AA), (0, HYDROPHILIC_AA), (1, POSITIVE_AA), (2, POLAR_AA), (3, OTHER_AA)):
    AA_MAP.update(**dict.fromkeys(sequence, index))


def map_aa_to_index(aa_seq: str):
    """This function converts amino acid sequence to index found in the `COLORS` dictionary."""
    aa_seq = aa_seq.replace(" ", "")  # remove all whitespace characters
    mapping = []
    for aa in aa_seq:
        mapping.append(AA_MAP[aa])
    return mapping
    
    
# This is your amino acid sequence
AA_SEQUENCE = \
    "ISQLGDESFR AICEYNSQAS DGVFKGPVSE TESFIKTQPG GHLAGYMGGR VDKDPRIGGM" \
    "WIVAHQVNLS AAATARKSLP KDEAKKLCSY KDILVVYEGV IKLATRALAN LEAFHNRPTL" \
    "TDIRGAVPGD TKAWELMPLR YMRLSELLSE ELHVAQNIDQ SFYIKNQHNC RDVEPEMDVT" \
    "LRIGLPQVFL FTSVKTFTNI"

# Adjust according to your style preferences
COLORS = {
    -1: "#FF0000",  # Negative
    0: "#000000",  # Hydrophobic
    1: "#0000FF",  # Positive,
    2: "#00FF00",  # Polar,
    3: "#FFFF00",  # Other - I think in my plot other was not included so you can just change this to another color
}


# convert your amino acid sequence to 
mapping = map_aa_to_index(AA_SEQUENCE)
# make sure this has two dimensions otherwise can't be plotted as an image
mapping = np.atleast_2d(mapping)

# create colormap based on your specified colors
cmap = matplotlib.colors.LinearSegmentedColormap.from_list("AA", list(COLORS.values()), N=5)
cmap.set_under((0, 0, 0, 0))
cmap.set_over((0, 0, 0, 0))

# figsize controls the ultimate size of the figure, although it's not absolute and is also controlled by the `aspect` parameter
plt.figure(figsize=(18, 3))
# aspect here corresponds to the aspect ratio between the vertical and horizontal axes
# Increasing this value make the figure slightly taller
plt.imshow(mapping, aspect=20, cmap=cmap)
plt.axis("off")