Contact
CoCalc Logo Icon
StoreFeaturesDocsShareSupport News AboutSign UpSign In
| Download
Views: 7339
1
"""This file contains code for use with "Think Stats",

2
by Allen B. Downey, available from greenteapress.com

3


4
Copyright 2014 Allen B. Downey

5
License: GNU GPLv3 http://www.gnu.org/licenses/gpl.html

6
"""

7


8
from __future__ import print_function, division

9


10
import unittest

11
import thinkdsp

12


13
import numpy as np

14


15


16
class Test(unittest.TestCase):

17
    def assertArrayAlmostEqual(self, a, b):

18
        total_error = np.sum(np.abs(a-b))

19
        self.assertAlmostEqual(total_error, 0)

20


21
    def testMakeSpectrum(self):

22
        # rfft with n even

23
        ys = np.arange(10)

24
        wave = thinkdsp.Wave(ys, framerate=10)

25
        spectrum = wave.make_spectrum()

26
        self.assertAlmostEqual(spectrum.hs[0], 45)

27
        wave2 = spectrum.make_wave()

28
        #print(wave2.ys)

29
        self.assertArrayAlmostEqual(wave.ys, wave2.ys)

30


31
        # rfft with n odd

32
        ys = np.arange(11)

33
        wave = thinkdsp.Wave(ys, framerate=10)

34
        spectrum = wave.make_spectrum()

35
        self.assertAlmostEqual(spectrum.hs[0], 55)

36


37
        # TODO: make rfft invertible when n is odd

38
        #wave2 = spectrum.make_wave()

39
        #print(wave2.ys)

40
        #self.assertArrayAlmostEqual(wave.ys, wave2.ys)

41


42
        # fft with n even

43
        ys = np.arange(10)

44
        wave = thinkdsp.Wave(ys, framerate=10)

45
        spectrum = wave.make_spectrum(full=True)

46
        self.assertAlmostEqual(spectrum.hs[0], 45)

47
        wave2 = spectrum.make_wave()

48
        #print(wave2.ys)

49
        self.assertArrayAlmostEqual(wave.ys, wave2.ys)

50


51
        # fft with n odd

52
        ys = np.arange(11)

53
        wave = thinkdsp.Wave(ys, framerate=10)

54
        spectrum = wave.make_spectrum(full=True)

55
        self.assertAlmostEqual(spectrum.hs[0], 55)

56
        wave2 = spectrum.make_wave()

57
        #print(wave2.ys)

58
        self.assertArrayAlmostEqual(wave.ys, wave2.ys)

59


60
    def testComplexSinusoid(self):

61
        signal = thinkdsp.ComplexSinusoid(440, 0.7, 1.1)

62
        result = signal.evaluate(2.1) * complex(-1.5, -0.5)

63
        self.assertAlmostEqual(result, -0.164353351475-1.09452637056j)

64


65
    def testChirp(self):

66
        signal = thinkdsp.Chirp(100, 200, 0.5)

67
        result = signal.evaluate([1.001, 1.002, 1.005])

68
        self.assertAlmostEqual(result[0], 0.5)

69
        self.assertAlmostEqual(result[2], -0.49384417)

70


71
    def testExpoChirp(self):

72
        signal = thinkdsp.ExpoChirp(100, 200, 0.5)

73
        result = signal.evaluate([0, 0.001, 0.002])

74
        self.assertAlmostEqual(result[0], 0.5)

75
        self.assertAlmostEqual(result[2], -0.27167286)

76


77
    def testWaveAdd(self):

78
        ys = np.array([1, 2, 3, 4])

79
        wave1 = thinkdsp.Wave(ys, framerate=1)

80
        wave2 = wave1.copy()

81
        wave2.shift(2)

82
        wave = wave1 + wave2

83


84
        self.assertEqual(len(wave), 6)

85
        self.assertAlmostEqual(sum(wave.ys), 20)

86


87
    def testDct(self):

88
        signal = thinkdsp.CosSignal(freq=2)

89
        wave = signal.make_wave(duration=1, framerate=8)

90
        dct = wave.make_dct()

91


92
        self.assertAlmostEqual(dct.fs[0], 0.25)

93


94
    def testImpulses(self):

95
        imp_sig = thinkdsp.Impulses([0.01, 0.4, 0.8, 1.2],

96
                                    amps=[1, 0.5, 0.25, 0.1])

97
        impulses = imp_sig.make_wave(start=0, duration=1.3,

98
                                     framerate=11025)

99


100
        self.assertEqual(len(impulses), 14332)

101


102


103
if __name__ == "__main__":

104
    unittest.main()

105


106