[Matplotlib] stuff I use often #viz

colors.py

# make a list of colours
cmap = colors.ListedColormap(['red','white','black'])

# define a range of numbers: -0.5-0.5, 0.5-1.5, 1.5-2.5
bounds=[-0.5, 0.5, 1.5, 2.5] 

# Generate a colormap index based on discrete intervals.
norm = colors.BoundaryNorm(bounds, cmap.N)

# plot
plt.imshow(im, cmap, norm)

line_animation.py

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from IPython.display import HTML, Image

# data
NUM = 100
x = np.random.randint(low=1, high=100, size=NUM)
y = np.random.randint(low=1, high=100, size=NUM)
x2 = np.random.randint(low=1, high=100, size=NUM)
y2 = np.random.randint(low=1, high=100, size=NUM)

# make the plot - useful to always use subplots to make it easy to add more
fig, ax = plt.subplots(figsize=(8,6))

# now draw the objects on the axis
# dunno why the trailing comma, some needed mpl trickery
line1, = ax.plot(x, y,   label=f"Line_1")
line2, = ax.plot(x2, y2, label=f"Line_2")

# specifiy all the labels
ax.set_title("da Lines")
ax.legend(loc="upper left")

# update the data of the objects defined above (line1, line2)
# could call a func to generate new y values, but here
# I'm just changing the window on the original array
def animate(num):
    line1.set_data(x[:num], y[:num])
    line2.set_data(x2[:num], y2[:num])
    return (line1, line2) # return multiple objects in a tuple

# declare a frames array if there are a lot of data
# else the animation can take forever
# the numbers in this array gets passed to the animate func above
frames = np.arange(0, len(x)+1000, 1000)

# don't need to pass in frames here, it defaults to drawing one frame per step 
ani = animation.FuncAnimation(
    fig, animate, frames, blit=True)

# to display an interactive animation in a jupyter notebook
# can get quite large and slow down jupyter, but it works well for smaller stuff.
HTML(ani.to_jshtml())

# so sometimes better to save to file and display as needed
# by default mpl tries to use ffmpeg so install that or investigate PillowWriter
ani.save('dynamic_images.mp4') # save as video
ani.save('dynamic_images.gif') # save as gif

wordcloud.py

# install wordcloud first: conda install -c conda-forge wordcloud

from wordcloud import WordCloud, STOPWORDS, ImageColorGenerator
#spacystopwords = spacy.lang.en.STOP_WORDS

def show_word_cloud(txt: str=txt, 
                    title: str="Fancy wordcloud title goes here"):
    """generates a wordcloud using https://github.com/amueller/word_cloud
    
    Args:
        txt: string of words
    Returns:
        fig: wordcloud as a matplotlib figure"""
    
    # make wordcloud
    wc = WordCloud(max_font_size=34, 
                   background_color="white",
                   stopwords=STOPWORDS)
	
    # set up the plot
    fig, ax = plt.subplots(figsize=(10,8))
    ax.axis("off")
    ax.set_title(f"{title}", fontsize=24)
    
    wordcloud = wc.generate(txt)
    ax.imshow(wordcloud, interpolation='bilinear')  # check interpolation
    
    plt.close()  # prevent displaying it in jupyter
    return fig