Snowfall Animation using matplotlib

Snowfall

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.animation import FuncAnimation

"""
Snowfall with python matplotlib
Read more here https://sukhbinder.wordpress.com/?p=5193

by Sukhbinder

"""

class Snowfall:
    def __init__(self, msg="", n_flakes=225, W=900, H=900):
        self.msg = msg
        self.nflakes = n_flakes
        self.W = W
        self.H = H
        self.COLORS = matplotlib.cm.bone(np.linspace(0, 1, self.nflakes))
        # Create new Figure and an Axes which fills it and set background black
        self.fig = plt.figure(figsize=(7, 7))
        self.ax = self.fig.add_axes([0, 0, 1, 1], frameon=False)
        self.ax.set_xlim(0, self.W)
        self.ax.set_xticks([])
        self.ax.set_ylim(0, self.H)
        self.ax.set_yticks([])
        self.fig.patch.set_facecolor('black')
        self.colors = self.COLORS.copy()
        self.positions = None
        self.fall = None
        self.step = None
        self.currstep = None
        self.sizes = None

        self.init()

        # Add text
        self.make_text(text=self.msg)

        # Construct the scatter which we will update during animation
        # as the snowflakes develop.
        self.scat = self.ax.scatter(self.position[:, 0], self.position[:, 1],
                                    s=self.sizes, lw=0.1, c=self.colors, edgecolors=self.colors,
                                    facecolors=self.colors)

    def make_text(self, text="Happy Holidays", color="Silver", fsize=25):
        X = self.W / 3
        Y = 100
        self.ax.text(X, Y, text, fontsize=fsize, color=color, alpha=0.5)
        for i in range(1, 15):
            plt.text(X, Y-(1.*i), text, fontsize=fsize +
                     (.05*i), alpha=0.03, color="white")

    def init(self):
        sfs = np.random.randint(2, size=self.nflakes)
        fall = np.random.rand(self.nflakes)
        step = np.random.rand(self.nflakes)
        no1 = sfs == 1

        fall[no1] = 2.0+fall[no1]*2.0
        step[no1] = 0.05+step[no1]*0.1

        fall[~no1] = 3.0+fall[~no1]*2.0
        step[~no1] = 0.05+step[~no1]*0.05

        xy = np.random.uniform(size=(self.nflakes, 2))

        xy = xy*self.W
        xy[:, 1] = self.H

        self.position = xy
        self.fall = fall
        self.step = step
        self.currstep = np.zeros(self.nflakes)
        self.sizes = np.random.uniform(1, 15, self.nflakes)

    def __call__(self, frame_number):
        # Make all colors more transparent as time progresses.
        self.colors[:, 3] -= 1.0/self.nflakes
        self.colors[:, 3] = np.clip(self.colors[:, 3], 0, 1)

        # move them
        dx = self.fall * np.cos(self.currstep)
        dy = self.fall
        self.position[:, 1] -= dy
        self.position[:, 0] -= dx
        self.currstep += self.step

        current_index = self.position[:, 1].argmin()
        xy = np.random.uniform(0, self.W, 2)
        xy[1] = self.H

        sfs_i = np.random.randint(2)
        if sfs_i == 1:
            fall_i = 2.0+np.random.rand()*2.0
            step_i = 0.05+np.random.rand()*0.1
        else:
            fall_i = 3.0+np.random.rand()*2.0
            step_i = 0.05+np.random.rand()*0.05

        self.position[current_index] = xy
        self.colors[current_index] = self.COLORS[np.random.choice(
            range(self.nflakes))]
        self.sizes[current_index] = np.random.uniform(1, 15, 1)
        self.currstep[current_index] = 0
        self.fall[current_index] = fall_i
        self.step[current_index] = step_i

        # Update the scatter collection, with the new colors, sizes and positions.
        self.scat.set_color(self.colors)
        self.scat.set_edgecolors(self.colors)
        self.scat.set_sizes(self.sizes)
        self.scat.set_offsets(self.position)

    def show(self):
        # Construct the animation, using the update function as the animation director.
        animation = FuncAnimation(self.fig, self.__call__, interval=75)
        plt.show()


if __name__ == "__main__":
    dd = Snowfall("Happy Holidays").show()