fourier_series_anime_saw.py

#!/usr/bin/env python

""" A Python script to produce fourier series animation of sawtooth wave."""

# Copyright (C) 2020 by Akira TAMAMORI

# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

# Commentary:
# You can see the original video from:
# https://twitter.com/i/status/1256034260848238594

# You can also find the original MATLAB script from:
# https://github.com/sh01k/teaching/blob/master/fourier_series_mov.m

# Code:


import math
import numpy
import numpy.matlib
import matplotlib.pyplot as plt
import matplotlib.animation as animation


# Period [s]
PERIOD = 1

# Sampling frequency for plot
SAMP_FREQ = 100

# Duration[s]
TIME_LEN = 4 * PERIOD  #
TIME_NUM = math.floor(TIME_LEN * SAMP_FREQ)  # 400

# Angle [rad]
ANGLE = numpy.linspace(0, 2 * math.pi, SAMP_FREQ)  # (100, 1)

# Order of Fourier Series
ORDER = numpy.array([1, 2, 3, 4])
ORDER_MAT = numpy.matlib.repmat(ORDER, TIME_NUM, 1)  # (400, 4)

fig = plt.figure(figsize=[8.0, 8.0])  # 800 x 800
ax1 = fig.add_axes([0.04, 0.85, 0.14, 0.14])
ax1.set_xlim(-0.6, 0.6)
ax1.set_ylim(-0.6, 0.6)

ax2 = fig.add_axes([0.24, 0.85, 0.74, 0.14])
ax2.set_ylim(-0.6, 0.6)

ax3 = fig.add_axes([0.04, 0.65, 0.14, 0.14])
ax3.set_xlim(-0.6, 0.6)
ax3.set_ylim(-0.6, 0.6)

ax4 = fig.add_axes([0.24, 0.65, 0.74, 0.14])
ax4.set_ylim(-0.6, 0.6)

ax5 = fig.add_axes([0.04, 0.45, 0.14, 0.14])
ax5.set_xlim(-0.6, 0.6)
ax5.set_ylim(-0.6, 0.6)

ax6 = fig.add_axes([0.24, 0.45, 0.74, 0.14])
ax6.set_ylim(-0.6, 0.6)

ax7 = fig.add_axes([0.04, 0.25, 0.14, 0.14])
ax7.set_xlim(-0.6, 0.6)
ax7.set_ylim(-0.6, 0.6)

ax8 = fig.add_axes([0.24, 0.25, 0.74, 0.14])
ax8.set_ylim(-0.6, 0.6)

ax9 = fig.add_axes([0.04, 0.05, 0.14, 0.14])
ax9.set_xlim(-0.6, 0.6)
ax9.set_ylim(-0.6, 0.6)

ax10 = fig.add_axes([0.24, 0.05, 0.74, 0.14])
ax10.set_ylim(-0.6, 0.6)

images = []

for t0 in range(TIME_NUM):

    # Time [s]
    time_axis = numpy.arange(0, TIME_NUM).T / SAMP_FREQ  # (400, )
    time_axis = time_axis[::-1]
    t = numpy.arange(t0, t0 + TIME_NUM).T / SAMP_FREQ  # (400, )
    t = numpy.expand_dims(t, axis=1)  # (400, 1)
    t_mat = numpy.matlib.repmat(t, 1, len(ORDER))  # (400, 4)

    # Fourier coefficients
    coef = -PERIOD / (math.pi * ORDER) * numpy.cos(math.pi * ORDER)  # Saw wave

    # Fourier bases
    phi = 2 * math.pi * ORDER_MAT * t_mat / PERIOD  # (400, 4)

    circ = numpy.array([coef * numpy.cos(phi[TIME_NUM-1, :]),
                        coef * numpy.sin(phi[TIME_NUM-1, :])])  # (2, 4, 4)

    sig = numpy.sum(numpy.matlib.repmat(
        coef, TIME_NUM, 1) * numpy.sin(phi), axis=1)

    # plot complex plane
    im = ax1.plot(coef[0] * numpy.cos(ANGLE),
                  coef[0] * numpy.sin(ANGLE),
                  color="k", linewidth=1.5)
    im += ax1.plot(coef[1] * numpy.cos(ANGLE) + circ[0, 0],
                   coef[1] * numpy.sin(ANGLE) + circ[1, 0],
                   color="k", linewidth=1.5)
    im += ax1.plot(coef[2] * numpy.cos(ANGLE) + circ[0, 0] + circ[0, 1],
                   coef[2] * numpy.sin(ANGLE) + circ[1, 0] + circ[1, 1],
                   color="k", linewidth=1.5)
    im += ax1.plot(coef[3] * numpy.cos(ANGLE) + circ[0, 0] + circ[0, 1] + circ[0, 2],
                   coef[3] * numpy.sin(ANGLE) + circ[1, 0] +
                   circ[1, 1] + circ[1, 2],
                   color="k", linewidth=1.5)
    im += ax1.plot([0, circ[0, 0]],
                   [0, circ[1, 0]],
                   color="b", linestyle="-", marker="o",
                   markerfacecolor="b", markersize=4)
    im += ax1.plot([circ[0, 0],
                    circ[0, 1] + circ[0, 0]],
                   [circ[1, 0],
                    circ[1, 1] + circ[1, 0]],
                   color="b", linestyle="-", marker="o",
                   markerfacecolor="b", markersize=4)
    im += ax1.plot([circ[0, 1] + circ[0, 0],
                    circ[0, 2] + circ[0, 1] + circ[0, 0]],
                   [circ[1, 1] + circ[1, 0],
                    circ[1, 2] + circ[1, 1] + circ[1, 0]],
                   color="b", linestyle="-", marker="o",
                   markerfacecolor="b", markersize=4)
    im += ax1.plot([circ[0, 2] + circ[0, 1] + circ[0, 0],
                    circ[0, 3] + circ[0, 2] + circ[0, 1] + circ[0, 0]],
                   [circ[1, 2] + circ[1, 1] + circ[1, 0],
                    circ[1, 3] + circ[1, 2] + circ[1, 1] + circ[1, 0]],
                   color="b", linestyle="-", marker="o",
                   markerfacecolor="b", markersize=4)
    im += ax1.plot([circ[0, 3] + circ[0, 2] + circ[0, 1] + circ[0, 0],
                    0.6],
                   [circ[1, 3] + circ[1, 2] + circ[1, 1] + circ[1, 0],
                    circ[1, 3] + circ[1, 2] + circ[1, 1] + circ[1, 0]],
                   color="b", linestyle=":", marker="o", markerfacecolor="b",
                   linewidth=1, markersize=4)

    # plot signal
    im += ax2.plot(time_axis, sig, linestyle="-", color="b", linewidth=1.5)

    # plot complex plane k=1
    im += ax3.plot(coef[0] * numpy.cos(ANGLE), coef[0] * numpy.sin(ANGLE),
                   color="k", linewidth=1.5)
    im += ax3.plot([0, circ[0, 0]], [0, circ[1, 0]],
                   linestyle="-", color="b", marker="o",
                   markerfacecolor="b", markersize=4)
    im += ax3.plot([circ[0, 0], 0.6], [circ[1, 0], circ[1, 0]],
                   linestyle=":", color="b", marker="o", markersize=4,
                   markerfacecolor="b", linewidth=1)

    # plot signal k=1
    im += ax4.plot(time_axis, coef[0] * numpy.sin(phi[:, 0]),
                   linestyle="-", color="b", linewidth=1.5)

    # plot complex plane k=2
    im += ax5.plot(coef[1] * numpy.cos(ANGLE), coef[1] * numpy.sin(ANGLE),
                   color="k", linewidth=1.5)
    im += ax5.plot([0, circ[0, 1]], [0, circ[1, 1]],
                   linestyle="-", color="b", marker="o",
                   markerfacecolor="b", markersize=4)
    im += ax5.plot([circ[0, 1], 0.6], [circ[1, 1], circ[1, 1]],
                   linestyle=":", color="b", marker="o", markersize=4,
                   markerfacecolor="b", linewidth=1)

    # plot signal k=2
    im += ax6.plot(time_axis, coef[1] * numpy.sin(phi[:, 1]),
                   linestyle="-", color="b", linewidth=1.5)

    # plot complex plane k=3
    im += ax7.plot(coef[2] * numpy.cos(ANGLE), coef[2] * numpy.sin(ANGLE),
                   color="k", linewidth=1.5)
    im += ax7.plot([0, circ[0, 2]], [0, circ[1, 2]],
                   linestyle="-", color="b", marker="o",
                   markerfacecolor="b", markersize=4)
    im += ax7.plot([circ[0, 2], 0.6], [circ[1, 2], circ[1, 2]],
                   linestyle=":", color="b", marker="o", markersize=4,
                   markerfacecolor="b", linewidth=1)

    # plot signal k=3
    im += ax8.plot(time_axis, coef[2] * numpy.sin(phi[:, 2]),
                   linestyle="-", color="b", linewidth=1.5)

    # plot complex plane k=4
    im += ax9.plot(coef[3] * numpy.cos(ANGLE), coef[3] * numpy.sin(ANGLE),
                   color="k", linewidth=1.5)
    im += ax9.plot([0, circ[0, 3]], [0, circ[1, 3]],
                   linestyle="-", color="b", marker="o",
                   markerfacecolor="b", markersize=4)
    im += ax9.plot([circ[0, 3], 0.6], [circ[1, 3], circ[1, 3]],
                   linestyle=":", color="b", marker="o", markersize=4,
                   markerfacecolor="b", linewidth=1)

    # plot signal k=4
    im += ax10.plot(time_axis, coef[3] * numpy.sin(phi[:, 3]),
                    linestyle="-", color="b", linewidth=1.5)

    images.append(im)

ANIME = animation.ArtistAnimation(fig, images, interval=40)
ANIME.save("plot_sawtooth.mp4", writer="ffmpeg", dpi=300)