tests.py 8.8 KB
Newer Older
1
2
3
import unittest
import numpy as np
import os
Gregory Ashton's avatar
Gregory Ashton committed
4
5
6
import shutil
import pyfstat

7

Gregory Ashton's avatar
Gregory Ashton committed
8
9
10
11
class Test(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        pass
12

Gregory Ashton's avatar
Gregory Ashton committed
13
14
15
16
17
18
    @classmethod
    def tearDownClass(cls):
        pass


class TestWriter(Test):
Gregory Ashton's avatar
Gregory Ashton committed
19
    label = "TestWriter"
20
21

    def test_make_cff(self):
Gregory Ashton's avatar
Gregory Ashton committed
22
        Writer = pyfstat.Writer(self.label, outdir=outdir)
23
        Writer.make_cff()
Gregory Ashton's avatar
Gregory Ashton committed
24
        self.assertTrue(os.path.isfile('./TestData/{}.cff'.format(self.label)))
25
26

    def test_run_makefakedata(self):
Gregory Ashton's avatar
Gregory Ashton committed
27
        Writer = pyfstat.Writer(self.label, outdir=outdir, duration=86400)
28
29
30
        Writer.make_cff()
        Writer.run_makefakedata()
        self.assertTrue(os.path.isfile(
Gregory Ashton's avatar
Gregory Ashton committed
31
            './TestData/H-48_H1_1800SFT_TestWriter-700000000-86400.sft'))
32
33

    def test_makefakedata_usecached(self):
Gregory Ashton's avatar
Gregory Ashton committed
34
        Writer = pyfstat.Writer(self.label, outdir=outdir, duration=86400)
35
36
        if os.path.isfile(Writer.sftfilepath):
            os.remove(Writer.sftfilepath)
Gregory Ashton's avatar
Gregory Ashton committed
37
        Writer.make_cff()
38
        Writer.run_makefakedata()
39
        time_first = os.path.getmtime(Writer.sftfilepath)
40
        Writer.run_makefakedata()
41
        time_second = os.path.getmtime(Writer.sftfilepath)
42
43
44
        self.assertTrue(time_first == time_second)
        os.system('touch {}'.format(Writer.config_file_name))
        Writer.run_makefakedata()
45
        time_third = os.path.getmtime(Writer.sftfilepath)
46
47
48
        self.assertFalse(time_first == time_third)


Gregory Ashton's avatar
Gregory Ashton committed
49
class TestBaseSearchClass(Test):
50
51
52
    def test_shift_matrix(self):
        BSC = pyfstat.BaseSearchClass()
        dT = 10
53
        a = BSC._shift_matrix(4, dT)
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
        b = np.array([[1, 2*np.pi*dT, 2*np.pi*dT**2/2.0, 2*np.pi*dT**3/6.0],
                      [0, 1, dT, dT**2/2.0],
                      [0, 0, 1, dT],
                      [0, 0, 0, 1]])
        self.assertTrue(np.array_equal(a, b))

    def test_shift_coefficients(self):
        BSC = pyfstat.BaseSearchClass()
        thetaA = np.array([10., 1e2, 10., 1e2])
        dT = 100

        # Calculate the 'long' way
        thetaB = np.zeros(len(thetaA))
        thetaB[3] = thetaA[3]
        thetaB[2] = thetaA[2] + thetaA[3]*dT
        thetaB[1] = thetaA[1] + thetaA[2]*dT + .5*thetaA[3]*dT**2
        thetaB[0] = thetaA[0] + 2*np.pi*(thetaA[1]*dT + .5*thetaA[2]*dT**2
                                         + thetaA[3]*dT**3 / 6.0)

        self.assertTrue(
            np.array_equal(
75
                thetaB, BSC._shift_coefficients(thetaA, dT)))
76
77
78
79
80

    def test_shift_coefficients_loop(self):
        BSC = pyfstat.BaseSearchClass()
        thetaA = np.array([10., 1e2, 10., 1e2])
        dT = 1e1
81
        thetaB = BSC._shift_coefficients(thetaA, dT)
82
83
        self.assertTrue(
            np.allclose(
84
                thetaA, BSC._shift_coefficients(thetaB, -dT),
85
86
87
                rtol=1e-9, atol=1e-9))


Gregory Ashton's avatar
Gregory Ashton committed
88
class TestComputeFstat(Test):
Gregory Ashton's avatar
Gregory Ashton committed
89
    label = "TestComputeFstat"
90
91

    def test_run_computefstatistic_single_point(self):
Gregory Ashton's avatar
Gregory Ashton committed
92
93
        Writer = pyfstat.Writer(self.label, outdir=outdir, duration=86400,
                                h0=1, sqrtSX=1)
94
95
96
        Writer.make_data()
        predicted_FS = Writer.predict_fstat()

97
98
        search = pyfstat.ComputeFstat(
            tref=Writer.tref,
99
            sftfilepattern='{}/*{}*sft'.format(Writer.outdir, Writer.label))
100
101
102
        FS = search.get_fullycoherent_twoF(
            Writer.tstart, Writer.tend, Writer.F0, Writer.F1, Writer.F2,
            Writer.Alpha, Writer.Delta)
103
        print predicted_FS, FS
Gregory Ashton's avatar
Gregory Ashton committed
104
        self.assertTrue(np.abs(predicted_FS-FS)/FS < 0.2)
105

106
    def run_computefstatistic_single_point_no_noise(self):
Gregory Ashton's avatar
Gregory Ashton committed
107
108
        Writer = pyfstat.Writer(self.label, outdir=outdir, add_noise=False,
                                duration=86400, h0=1, sqrtSX=1)
109
110
111
112
113
        Writer.make_data()
        predicted_FS = Writer.predict_fstat()

        search = pyfstat.ComputeFstat(
            tref=Writer.tref, assumeSqrtSX=1,
114
            sftfilepattern='{}/*{}*sft'.format(Writer.outdir, Writer.label))
115
116
117
        FS = search.get_fullycoherent_twoF(
            Writer.tstart, Writer.tend, Writer.F0, Writer.F1, Writer.F2,
            Writer.Alpha, Writer.Delta)
118
119
        print predicted_FS, FS
        self.assertTrue(np.abs(predicted_FS-FS)/FS < 0.2)
Gregory Ashton's avatar
Gregory Ashton committed
120

121
    def test_injectSources_from_file(self):
Gregory Ashton's avatar
Gregory Ashton committed
122
123
        Writer = pyfstat.Writer(self.label, outdir=outdir, add_noise=False,
                                duration=86400, h0=1, sqrtSX=1)
124
125
126
127
128
129
130
131
        Writer.make_cff()
        injectSources = Writer.config_file_name

        search = pyfstat.ComputeFstat(
            tref=Writer.tref, assumeSqrtSX=1, injectSources=injectSources,
            minCoverFreq=28, maxCoverFreq=32, minStartTime=Writer.tstart,
            maxStartTime=Writer.tstart+Writer.duration,
            detectors=Writer.detectors)
132
133
134
        FS = search.get_fullycoherent_twoF(
            Writer.tstart, Writer.tend, Writer.F0, Writer.F1, Writer.F2,
            Writer.Alpha, Writer.Delta)
135
136
137
138
139
140
        Writer.make_data()
        predicted_FS = Writer.predict_fstat()
        print predicted_FS, FS
        self.assertTrue(np.abs(predicted_FS-FS)/FS < 0.2)


Gregory Ashton's avatar
Gregory Ashton committed
141
class TestSemiCoherentGlitchSearch(Test):
Gregory Ashton's avatar
Gregory Ashton committed
142
    label = "TestSemiCoherentGlitchSearch"
Gregory Ashton's avatar
Gregory Ashton committed
143

144
    def test_get_semicoherent_nglitch_twoF(self):
Gregory Ashton's avatar
Gregory Ashton committed
145
146
        duration = 10*86400
        dtglitch = .5*duration
147
        delta_F0 = 0
Gregory Ashton's avatar
Gregory Ashton committed
148
149
        h0 = 1
        sqrtSX = 1
150
151
        Writer = pyfstat.GlitchWriter(
            self.label, outdir=outdir, duration=duration, dtglitch=dtglitch,
Gregory Ashton's avatar
Gregory Ashton committed
152
            delta_F0=delta_F0, sqrtSX=sqrtSX, h0=h0)
153
154
155
156

        Writer.make_data()

        search = pyfstat.SemiCoherentGlitchSearch(
157
            label=self.label, outdir=outdir,
158
            sftfilepattern='{}/*{}*sft'.format(Writer.outdir, Writer.label),
159
160
            tref=Writer.tref, minStartTime=Writer.tstart,
            maxStartTime=Writer.tend, nglitch=1)
161

162
163
164
        FS = search.get_semicoherent_nglitch_twoF(
            Writer.F0, Writer.F1, Writer.F2, Writer.Alpha, Writer.Delta,
            Writer.delta_F0, Writer.delta_F1, search.minStartTime+dtglitch)
165
166

        # Compute the predicted semi-coherent glitch Fstat
167
168
        minStartTime = Writer.tstart
        maxStartTime = Writer.tend
169

170
        Writer.maxStartTime = minStartTime + dtglitch
171
172
        FSA = Writer.predict_fstat()

173
174
        Writer.tstart = minStartTime + dtglitch
        Writer.tend = maxStartTime
175
176
        FSB = Writer.predict_fstat()

177
        print FSA, FSB
178
179
180
        predicted_FS = (FSA + FSB)

        print(predicted_FS, FS)
Gregory Ashton's avatar
Gregory Ashton committed
181
        self.assertTrue(np.abs((FS - predicted_FS))/predicted_FS < 0.3)
182
183


Gregory Ashton's avatar
Gregory Ashton committed
184
class TestMCMCSearch(Test):
Gregory Ashton's avatar
Gregory Ashton committed
185
    label = "TestMCMCSearch"
186
187

    def test_fully_coherent(self):
Gregory Ashton's avatar
Gregory Ashton committed
188
189
        h0 = 1
        sqrtSX = 1
190
191
192
        F0 = 30
        F1 = -1e-10
        F2 = 0
193
        minStartTime = 700000000
Gregory Ashton's avatar
Gregory Ashton committed
194
        duration = 1 * 86400
195
        maxStartTime = minStartTime + duration
196
197
        Alpha = 5e-3
        Delta = 1.2
198
        tref = minStartTime
199
200
        Writer = pyfstat.Writer(F0=F0, F1=F1, F2=F2, label=self.label,
                                h0=h0, sqrtSX=sqrtSX,
201
                                outdir=outdir, tstart=minStartTime,
202
                                Alpha=Alpha, Delta=Delta, tref=tref,
203
                                duration=duration,
204
                                Band=4)
205
206
207
208

        Writer.make_data()
        predicted_FS = Writer.predict_fstat()

Gregory Ashton's avatar
Gregory Ashton committed
209
210
        theta = {'F0': {'type': 'norm', 'loc': F0, 'scale': np.abs(1e-10*F0)},
                 'F1': {'type': 'norm', 'loc': F1, 'scale': np.abs(1e-10*F1)},
211
212
                 'F2': F2, 'Alpha': Alpha, 'Delta': Delta}

Gregory Ashton's avatar
Gregory Ashton committed
213
        search = pyfstat.MCMCSearch(
Gregory Ashton's avatar
Gregory Ashton committed
214
            label=self.label, outdir=outdir, theta_prior=theta, tref=tref,
215
            sftfilepattern='{}/*{}*sft'.format(Writer.outdir, Writer.label),
216
            minStartTime=minStartTime, maxStartTime=maxStartTime,
217
            nsteps=[100, 100], nwalkers=100, ntemps=2, log10beta_min=-1)
Gregory Ashton's avatar
Gregory Ashton committed
218
        search.setup_burnin_convergence_testing()
219
        search.run(create_plots=False)
220
221
222
223
        _, FS = search.get_max_twoF()

        print('Predicted twoF is {} while recovered is {}'.format(
                predicted_FS, FS))
Gregory Ashton's avatar
Gregory Ashton committed
224
225
        self.assertTrue(
            FS > predicted_FS or np.abs((FS-predicted_FS))/predicted_FS < 0.3)
226

227
    def test_multi_stage(self):
Gregory Ashton's avatar
Gregory Ashton committed
228
        Writer = pyfstat.Writer(F0=10, duration=86400, h0=1, sqrtSX=1)
229
230
231
232
233
234
235
236
237
238
239
240
241
        Writer.make_cff()

        theta = {'F0': {'type': 'norm', 'loc': 10, 'scale': 1e-2},
                 'F1': 0, 'F2': 0, 'Alpha': 0, 'Delta': 0}

        search = pyfstat.MCMCSearch(
            label=self.label, outdir=outdir, theta_prior=theta,
            tref=Writer.tref, injectSources=Writer.config_file_name,
            minStartTime=Writer.minStartTime, maxStartTime=Writer.maxStartTime,
            nsteps=[5, 5], nwalkers=20, ntemps=1, detectors='H1',
            minCoverFreq=9, maxCoverFreq=11)
        search.run(create_plots=False)

242
243

if __name__ == '__main__':
Gregory Ashton's avatar
Gregory Ashton committed
244
245
246
    outdir = 'TestData'
    if os.path.isdir(outdir):
        shutil.rmtree(outdir)
247
    unittest.main()
Gregory Ashton's avatar
Gregory Ashton committed
248
249
250
    if os.path.isdir(outdir):
        shutil.rmtree(outdir)