Difference between revisions of "How To set a Python matplotlib axis to display from -Pi to Pi"
(Created page with "[https://www.python.org/ Python] and [https://matplotlib.org/ Matplotlib] are great tools for doing Science! In particular, I like to use the [https://www.spyder-ide.org/ Spyd...") |
(No difference)
|
Latest revision as of 11:06, 8 May 2020
Python and Matplotlib are great tools for doing Science! In particular, I like to use the Spyder IDE. Matplotlib makes it easy to create graphs from datasets. In the world of DSP, an example of this usage is plotting sampled data sets in the time domain and frequency domain. However, by default your frequency domain plot's x-axis will simply be labelled with the indices of your sampled data array. In DSP it is often desirable to remap this range of indices (0 to N-1) to the frequencies of -Pi to Pi. Unfortunately there isn't a 1 step solution to this remapping. However, with a little bit of code you can accomplish this relatively easily.
The code below with show you how to go from this most basic graph of a DFT:
to the much, prettier, shifted and relabeled graph of a DFT
To remap your DFT so that it is centered about 0 frequency you can use the following dft_shift function.
1 import matplotlib
2 import matplotlib.pyplot as plt
3 import numpy as np
4
5 # define dft_shift to center the DFT about 0. Make sure the final length is odd so 0 is at the center
6 def dft_shift(X):
7 N = len(X)
8 if (N % 2 == 0):
9 # even-length: return N+1 values
10 return np.arange(-int(N/2), int(N/2) + 1), np.concatenate((X[int(N/2):], X[:int(N/2)+1]))
11 else:
12 # odd-length: return N values
13 return np.arange(-int((N-1)/2), int((N-1)/2) + 1), np.concatenate((X[int((N+1)/2):], X[:int((N+1)/2)]))
To change a matplotlib plot's x-axis to display from -Pi to Pi you can use the following lines of code. X is assumed to be your sampled data array
1 n_pi = np.linspace(-np.pi,np.pi,len(X)) # create array from -pi to pi that is the same length as your sampled data array
2 plt.xlim([-np.pi, np.pi]) # set x-axis display range from -Pi to Pi
3 plt.xticks([-np.pi, -np.pi/2, 0, np.pi/2, np.pi],
4 [r'$-\pi$', r'$-\pi/2$', r'$0$', r'$+\pi/2$', r'$\pi$']) # display ticks in pi format rather than 3.14 format
5 plt.stem(n_pi, abs(X)) # plot sampled data array X using n_pi as the indicies
The complete Python script below puts it all together and walks you through graphing the same data in various formats.
1 import matplotlib
2 import matplotlib.pyplot as plt
3 import numpy as np
4
5 # define dft_shift to center the DFT about 0. Make sure the final length is odd so 0 is at the center
6 def dft_shift(X):
7 N = len(X)
8 if (N % 2 == 0):
9 # even-length: return N+1 values
10 return np.arange(-int(N/2), int(N/2) + 1), np.concatenate((X[int(N/2):], X[:int(N/2)+1]))
11 else:
12 # odd-length: return N values
13 return np.arange(-int((N-1)/2), int((N-1)/2) + 1), np.concatenate((X[int((N+1)/2):], X[:int((N+1)/2)]))
14
15 # define the size of the figures
16 plt.rcParams["figure.figsize"] = (14,4)
17
18 # create simple baseband signal
19 N = 64
20 x = np.ones(N)
21
22 # plot x
23 plt.figure()
24 plt.title("baseband signal of 1s")
25 plt.stem(x)
26
27 # plot DFT(x) = X
28 X = np.fft.fft(x)
29 plt.figure()
30 plt.title("basic DFT")
31 plt.stem(abs(X))
32
33 # plot shifted X with default indicies
34 n,X_shift = dft_shift(X) # get the shifted version of X. n contains integer indicies centered about 0
35 plt.figure()
36 plt.title("shifted DFT with default indicies")
37 plt.stem(abs(X_shift)) # plot shifted X with "default" 0 to len(X_shift) indicies
38
39 # plot shifted X with n as indicies
40 n,X_shift = dft_shift(X) # get the shifted version of X. n contains integer indicies centered about 0
41 plt.figure()
42 plt.title("shifted DFT with n as indicies")
43 plt.stem(n, abs(X_shift)) # plot shifted X using n as indicies
44
45 # plot shifted X with plot scaled from -Pi to Pi (but still with integer indicies)
46 n,X_shift = dft_shift(X) # get the shifted version of X. n contain indicies centered about 0
47 plt.figure()
48 plt.title("shifted DFT from -Pi to Pi with n as indicies (notice only 7 points are plotted)")
49 plt.xlim([-np.pi, np.pi]) # x-axis display range from -Pi to Pi
50 plt.xticks([-np.pi, -np.pi/2, 0, np.pi/2, np.pi],
51 [r'$-\pi$', r'$-\pi/2$', r'$0$', r'$+\pi/2$', r'$\pi$']) # display ticks in pi format rather than 3.14 format
52 plt.stem(n, abs(X_shift)) # plot shifted X using n as the indicies
53
54 # plot shifted X with X-axis indicies from -Pi to Pi and plot scaled from -Po to Pi
55 n,X_shift = dft_shift(X) # get the shifted version of X. n contain indicies centered about 0
56 n_pi = np.linspace(-np.pi,np.pi,len(n)) # create array from -pi to pi that is the same length as n
57 plt.figure()
58 plt.title("shifted DFT from -Pi to Pi with n_pi as indicies")
59 plt.xlim([-np.pi, np.pi]) # x-axis display range from -Pi to Pi
60 plt.xticks([-np.pi, -np.pi/2, 0, np.pi/2, np.pi],
61 [r'$-\pi$', r'$-\pi/2$', r'$0$', r'$+\pi/2$', r'$\pi$']) # display ticks in pi format rather than 3.14 format
62 plt.stem(n_pi, abs(X_shift)) # plot shifted X using n_pi as the indicies
