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.spyderide.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 xaxis 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 N1) 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 # evenlength: 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 # oddlength: return N values
13 return np.arange(int((N1)/2), int((N1)/2) + 1), np.concatenate((X[int((N+1)/2):], X[:int((N+1)/2)]))
To change a matplotlib plot's xaxis 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 xaxis 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 # evenlength: 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 # oddlength: return N values
13 return np.arange(int((N1)/2), int((N1)/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]) # xaxis 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 Xaxis 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]) # xaxis 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