create_injection_test.py 7.84 KB
Newer Older
1
2
3
4
import os
import shutil
import unittest

Gregory Ashton's avatar
Gregory Ashton committed
5
import numpy as np
6
7

import bilby_pipe
Gregory Ashton's avatar
Gregory Ashton committed
8
9
from bilby_pipe.input import Input
from bilby_pipe.utils import BilbyPipeError, parse_args
10
11
12


class TestParser(unittest.TestCase):
Gregory Ashton's avatar
Gregory Ashton committed
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
    def test_parser_defaults(self):
        example_prior_file = "tests/example_prior.prior"

        known_args_list = [example_prior_file, "-n", "1"]
        parser = bilby_pipe.create_injections.create_parser()
        args, unknown_args = parse_args(known_args_list, parser)

        self.assertEqual(args.prior_file, example_prior_file)
        self.assertEqual(args.n_injection, 1)
        self.assertEqual(args.extension, "dat")

    def test_parser_with_prior_file(self):
        example_prior_file = "tests/example_prior.prior"

        known_args_list = [example_prior_file, "--n-injection", "3", "-s", "1234"]
        parser = bilby_pipe.create_injections.create_parser()
        args, unknown_args = parse_args(known_args_list, parser)

        self.assertEqual(args.prior_file, example_prior_file)
        self.assertEqual(args.n_injection, 3)
        self.assertEqual(args.generation_seed, 1234)

    def test_parser_with_default_prior_file(self):
        known_args_list = ["4s", "--n-injection", "3"]
        parser = bilby_pipe.create_injections.create_parser()
        args, unknown_args = parse_args(known_args_list, parser)

        self.assertEqual(args.prior_file, "4s")
        self.assertEqual(args.n_injection, 3)

    def test_parser_with_json(self):
        example_prior_file = "tests/example_prior.prior"

        known_args_list = [example_prior_file, "--n-injection", "3", "-e", "json"]
        parser = bilby_pipe.create_injections.create_parser()
        args, unknown_args = parse_args(known_args_list, parser)

        self.assertEqual(args.extension, "json")


class TestCreateInjections(unittest.TestCase):
54
    def setUp(self):
55
56
        self.outdir = "outdir"
        self.example_prior_file = "tests/example_prior.prior"
57
        self.filename = f"{self.outdir}/injection.dat"
58
59
60
61
62
63

    def tearDown(self):
        try:
            shutil.rmtree(self.outdir)
        except FileNotFoundError:
            pass
Gregory Ashton's avatar
Gregory Ashton committed
64
65

    def test_create_injection_file(self):
66
        filename = f"{self.outdir}/injections.dat"
Gregory Ashton's avatar
Gregory Ashton committed
67
68
69
        prior_file = self.example_prior_file
        n_injection = 3
        bilby_pipe.create_injections.create_injection_file(
70
71
72
73
74
            filename,
            n_injection,
            prior_file=prior_file,
            generation_seed=None,
            extension="dat",
Gregory Ashton's avatar
Gregory Ashton committed
75
76
77
78
79
80
        )
        self.assertTrue(os.path.isfile(filename))
        injections = np.genfromtxt(filename, names=True)
        self.assertEqual(len(injections), n_injection)

    def test_create_injection_file_dat_ext(self):
81
        filename = f"{self.outdir}/injections"
Gregory Ashton's avatar
Gregory Ashton committed
82
83
84
        prior_file = self.example_prior_file
        n_injection = 3
        bilby_pipe.create_injections.create_injection_file(
85
86
87
88
89
            filename,
            n_injection,
            prior_file=prior_file,
            generation_seed=None,
            extension="dat",
Gregory Ashton's avatar
Gregory Ashton committed
90
91
92
93
94
95
96
        )
        actual_filename = filename + ".dat"
        self.assertTrue(os.path.isfile(actual_filename))
        df = Input.read_dat_injection_file(actual_filename)
        self.assertEqual(len(df), n_injection)

    def test_create_injection_file_json_ext(self):
97
        filename = f"{self.outdir}/injections"
Gregory Ashton's avatar
Gregory Ashton committed
98
99
100
        prior_file = self.example_prior_file
        n_injection = 3
        bilby_pipe.create_injections.create_injection_file(
101
102
103
104
105
            filename,
            n_injection,
            prior_file=prior_file,
            generation_seed=None,
            extension="json",
Gregory Ashton's avatar
Gregory Ashton committed
106
107
108
109
110
111
        )
        actual_filename = filename + ".json"
        self.assertTrue(os.path.isfile(actual_filename))
        df = Input.read_json_injection_file(actual_filename)
        self.assertEqual(len(df), n_injection)

112
    def test_create_injection_file_with_gps_file(self):
113
        filename = f"{self.outdir}/injections"
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
        prior_file = self.example_prior_file
        n_injection = 2
        bilby_pipe.create_injections.create_injection_file(
            filename,
            n_injection,
            prior_file=prior_file,
            generation_seed=None,
            extension="json",
            gps_file="tests/gps_file.txt",
        )
        filename += ".json"
        gps_vals = np.loadtxt("tests/gps_file.txt")

        df = Input.read_json_injection_file(filename)
        self.assertEqual(len(df), n_injection)
        self.assertTrue(
            len(df.columns.values) > 1, f"Column names: {df.columns.values}"
        )
        self.assertAlmostEqual(
Gregory Ashton's avatar
Gregory Ashton committed
133
            df["geocent_time"].iloc[0] / 100, gps_vals[0] / 100, places=1
134
135
        )

Gregory Ashton's avatar
Gregory Ashton committed
136
    def test_create_injection_file_json(self):
137
        filename = f"{self.outdir}/injections.json"
Gregory Ashton's avatar
Gregory Ashton committed
138
139
140
        prior_file = self.example_prior_file
        n_injection = 3
        bilby_pipe.create_injections.create_injection_file(
141
142
143
144
145
            filename,
            n_injection,
            prior_file=prior_file,
            generation_seed=None,
            extension="dat",
Gregory Ashton's avatar
Gregory Ashton committed
146
147
148
149
150
151
152
        )
        actual_filename = filename
        self.assertTrue(os.path.isfile(actual_filename))
        df = Input.read_json_injection_file(actual_filename)
        self.assertEqual(len(df), n_injection)

    def test_create_injection_file_generation_seed(self):
153
        filename = f"{self.outdir}/injections"
Gregory Ashton's avatar
Gregory Ashton committed
154
155
156
        prior_file = self.example_prior_file
        n_injection = 3
        bilby_pipe.create_injections.create_injection_file(
157
            filename + "_A", n_injection, prior_file=prior_file, generation_seed=123
Gregory Ashton's avatar
Gregory Ashton committed
158
159
160
161
        )
        injectionsA = np.genfromtxt(filename + "_A.dat", names=True)

        bilby_pipe.create_injections.create_injection_file(
162
            filename + "_B", n_injection, prior_file=prior_file, generation_seed=123
Gregory Ashton's avatar
Gregory Ashton committed
163
164
165
166
        )
        injectionsB = np.genfromtxt(filename + "_B.dat", names=True)

        bilby_pipe.create_injections.create_injection_file(
167
            filename + "_C", n_injection, prior_file=prior_file, generation_seed=321
Gregory Ashton's avatar
Gregory Ashton committed
168
169
170
171
172
173
174
175
176
177
        )
        injectionsC = np.genfromtxt(filename + "_C.dat", names=True)

        self.assertTrue(np.all(injectionsA == injectionsB))
        self.assertFalse(np.all(injectionsA == injectionsC))

    def test_n_injection_error(self):
        with self.assertRaises(BilbyPipeError):
            n_injection = None
            bilby_pipe.create_injections.create_injection_file(
178
                self.filename, n_injection, prior_file=self.example_prior_file
Gregory Ashton's avatar
Gregory Ashton committed
179
180
            )

181
        with self.assertRaises(BilbyPipeError):
Gregory Ashton's avatar
Gregory Ashton committed
182
183
            n_injection = -1
            bilby_pipe.create_injections.create_injection_file(
184
                self.filename, n_injection, prior_file=self.example_prior_file
Gregory Ashton's avatar
Gregory Ashton committed
185
186
187
188
189
            )

        with self.assertRaises(BilbyPipeError):
            n_injection = np.inf
            bilby_pipe.create_injections.create_injection_file(
190
                self.filename, n_injection, prior_file=self.example_prior_file
Gregory Ashton's avatar
Gregory Ashton committed
191
            )
192
193

    def test_unknown_prior_file(self):
Gregory Ashton's avatar
Gregory Ashton committed
194
        prior_file = "not_a_file"
195
        with self.assertRaises(FileNotFoundError):
Gregory Ashton's avatar
Gregory Ashton committed
196
            bilby_pipe.create_injections.create_injection_file(
197
                self.filename, 1, prior_file=prior_file
Gregory Ashton's avatar
Gregory Ashton committed
198
199
200
201
202
203
            )

    def test_none_prior_file(self):
        prior_file = None
        with self.assertRaises(BilbyPipeError):
            bilby_pipe.create_injections.create_injection_file(
204
                self.filename, 1, prior_file=prior_file
Gregory Ashton's avatar
Gregory Ashton committed
205
206
207
208
209
            )

    def test_unknown_ext(self):
        with self.assertRaises(BilbyPipeError):
            bilby_pipe.create_injections.create_injection_file(
210
                "test", 1, self.example_prior_file, extension="other"
Gregory Ashton's avatar
Gregory Ashton committed
211
212
213
214
215
            )

    def test_unknown_ext_from_filename(self):
        with self.assertRaises(BilbyPipeError):
            bilby_pipe.create_injections.create_injection_file(
216
                "test.other", 1, self.example_prior_file
Gregory Ashton's avatar
Gregory Ashton committed
217
            )
218
219


220
if __name__ == "__main__":
221
    unittest.main()