main_test.py 8.67 KB
Newer Older
Gregory Ashton's avatar
Gregory Ashton committed
1
import copy
2
import os
Gregory Ashton's avatar
Gregory Ashton committed
3
import shutil
4
import unittest
Gregory Ashton's avatar
Gregory Ashton committed
5

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

Gregory Ashton's avatar
Gregory Ashton committed
8
import bilby_pipe
Gregory Ashton's avatar
Gregory Ashton committed
9
from bilby_pipe.utils import BilbyPipeError
Gregory Ashton's avatar
Gregory Ashton committed
10

Gregory Ashton's avatar
Gregory Ashton committed
11

Gregory Ashton's avatar
Gregory Ashton committed
12 13 14
class TestMainInput(unittest.TestCase):
    def setUp(self):
        self.directory = os.path.abspath(os.path.dirname(__file__))
Gregory Ashton's avatar
Gregory Ashton committed
15
        self.outdir = "outdir"
Gregory Ashton's avatar
Gregory Ashton committed
16
        self.known_args_list = [
Gregory Ashton's avatar
Gregory Ashton committed
17 18 19 20 21 22
            "tests/test_main_input.ini",
            "--submit",
            "--outdir",
            self.outdir,
        ]
        self.unknown_args_list = ["--argument", "value"]
Gregory Ashton's avatar
Gregory Ashton committed
23
        self.all_args_list = self.known_args_list + self.unknown_args_list
Gregory Ashton's avatar
Gregory Ashton committed
24
        self.parser = bilby_pipe.main.create_parser()
Gregory Ashton's avatar
Gregory Ashton committed
25 26
        self.args = self.parser.parse_args(self.known_args_list)
        self.inputs = bilby_pipe.main.MainInput(
Gregory Ashton's avatar
Gregory Ashton committed
27 28
            *self.parser.parse_known_args(self.all_args_list)
        )
Gregory Ashton's avatar
Gregory Ashton committed
29

Gregory Ashton's avatar
Gregory Ashton committed
30
        self.test_gps_file = "tests/gps_file.txt"
31

Gregory Ashton's avatar
Gregory Ashton committed
32 33 34 35 36 37
    def tearDown(self):
        shutil.rmtree(self.outdir)
        del self.args
        del self.inputs

    def test_ini(self):
Gregory Ashton's avatar
Gregory Ashton committed
38
        self.assertEqual(self.inputs.ini, os.path.relpath(self.args.ini))
Gregory Ashton's avatar
Gregory Ashton committed
39

Gregory Ashton's avatar
Gregory Ashton committed
40 41 42
    def test_initialdir(self):
        self.assertEqual(self.inputs.initialdir, os.getcwd())

Gregory Ashton's avatar
Gregory Ashton committed
43
    def test_ini_not_a_file(self):
44
        with self.assertRaises(FileNotFoundError):
Gregory Ashton's avatar
Gregory Ashton committed
45
            self.inputs.ini = "not_a_file"
Gregory Ashton's avatar
Gregory Ashton committed
46

Gregory Ashton's avatar
Gregory Ashton committed
47 48 49
    def test_submit(self):
        self.assertEqual(self.inputs.submit, self.args.submit)

Gregory Ashton's avatar
Gregory Ashton committed
50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77
    def test_request_cpus(self):
        self.assertEqual(self.inputs.request_cpus, self.args.request_cpus)

    def test_request_memory(self):
        memory = "{} GB".format(self.args.request_memory)
        self.assertEqual(self.inputs.request_memory, memory)

    def test_request_memory_generation_default_non_roq(self):
        mem_int = bilby_pipe.utils.request_memory_generation_lookup(
            self.args.duration, roq=False
        )
        memory = "{} GB".format(mem_int)
        self.assertEqual(self.inputs.request_memory_generation, memory)

    def test_request_memory_generation_default_roq(self):
        self.inputs.likelihood_type = "ROQGravitationalWaveTransient"
        mem_int = bilby_pipe.utils.request_memory_generation_lookup(
            self.args.duration, roq=True
        )
        memory = "{} GB".format(mem_int)
        self.assertEqual(self.inputs.request_memory_generation, memory)

    def test_request_memory_generation_set(self):
        self.args.request_memory_generation = 4
        inputs = bilby_pipe.main.MainInput(self.args, [])
        memory = "{} GB".format(self.args.request_memory_generation)
        self.assertEqual(inputs.request_memory_generation, memory)

Avi Vajpeyi's avatar
Avi Vajpeyi committed
78 79 80 81 82 83 84 85 86 87
    def test_notification_set(self):
        self.args.notification = "Always"
        inputs = bilby_pipe.main.MainInput(self.args, [])
        self.assertEqual(inputs.notification, "Always")

    def test_notification_error_riased_set(self):
        self.args.notification = "Sometimes"
        with self.assertRaises(BilbyPipeError):
            bilby_pipe.main.MainInput(self.args, [])

Gregory Ashton's avatar
Gregory Ashton committed
88 89 90 91
    def test_label(self):
        self.assertEqual(self.inputs.label, self.args.label)

    def test_coherence_test(self):
Gregory Ashton's avatar
Gregory Ashton committed
92
        self.assertEqual(self.inputs.coherence_test, self.args.coherence_test)
Gregory Ashton's avatar
Gregory Ashton committed
93 94 95 96 97 98

    def test_accounting(self):
        self.assertEqual(self.inputs.accounting, self.args.accounting)

    def test_detectors_single(self):
        # Test the detector set in the ini file
Gregory Ashton's avatar
Gregory Ashton committed
99
        self.assertEqual(self.inputs.detectors, ["H1"])
Gregory Ashton's avatar
Gregory Ashton committed
100 101 102

        # Test setting a single detector directly in the args as a string
        args = copy.copy(self.args)
Gregory Ashton's avatar
Gregory Ashton committed
103 104 105
        args.detectors = "L1"
        inputs = bilby_pipe.main.MainInput(args, [])
        self.assertEqual(inputs.detectors, ["L1"])
Gregory Ashton's avatar
Gregory Ashton committed
106

Gregory Ashton's avatar
Gregory Ashton committed
107
        args.detectors = ["L1"]
Gregory Ashton's avatar
Gregory Ashton committed
108
        inputs = bilby_pipe.main.MainInput(args, self.unknown_args_list)
Gregory Ashton's avatar
Gregory Ashton committed
109
        self.assertEqual(inputs.detectors, ["L1"])
Gregory Ashton's avatar
Gregory Ashton committed
110

111
        with self.assertRaises(BilbyPipeError):
Gregory Ashton's avatar
Gregory Ashton committed
112
            args.detectors = "A1"
Gregory Ashton's avatar
Gregory Ashton committed
113 114
            inputs = bilby_pipe.main.MainInput(args, self.unknown_args_list)

115
        with self.assertRaises(BilbyPipeError):
Gregory Ashton's avatar
Gregory Ashton committed
116 117 118 119 120
            args.detectors = None
            inputs = bilby_pipe.main.MainInput(args, self.unknown_args_list)

    def test_detectors_multiple(self):
        args = copy.copy(self.args)
Gregory Ashton's avatar
Gregory Ashton committed
121
        args.detectors = "H1 L1"
Gregory Ashton's avatar
Gregory Ashton committed
122
        inputs = bilby_pipe.main.MainInput(args, self.unknown_args_list)
Gregory Ashton's avatar
Gregory Ashton committed
123
        self.assertEqual(inputs.detectors, ["H1", "L1"])
Gregory Ashton's avatar
Gregory Ashton committed
124

Gregory Ashton's avatar
Gregory Ashton committed
125
        args.detectors = "L1 H1"
Gregory Ashton's avatar
Gregory Ashton committed
126
        inputs = bilby_pipe.main.MainInput(args, self.unknown_args_list)
Gregory Ashton's avatar
Gregory Ashton committed
127
        self.assertEqual(inputs.detectors, ["H1", "L1"])
Gregory Ashton's avatar
Gregory Ashton committed
128

Gregory Ashton's avatar
Gregory Ashton committed
129
        args.detectors = ["L1 H1"]
Gregory Ashton's avatar
Gregory Ashton committed
130
        inputs = bilby_pipe.main.MainInput(args, self.unknown_args_list)
Gregory Ashton's avatar
Gregory Ashton committed
131
        self.assertEqual(inputs.detectors, ["H1", "L1"])
Gregory Ashton's avatar
Gregory Ashton committed
132

Gregory Ashton's avatar
Gregory Ashton committed
133
        args.detectors = ["L1", "H1"]
Gregory Ashton's avatar
Gregory Ashton committed
134
        inputs = bilby_pipe.main.MainInput(args, self.unknown_args_list)
Gregory Ashton's avatar
Gregory Ashton committed
135
        self.assertEqual(inputs.detectors, ["H1", "L1"])
Gregory Ashton's avatar
Gregory Ashton committed
136

Gregory Ashton's avatar
Gregory Ashton committed
137
        args.detectors = ["H1", "L1"]
Gregory Ashton's avatar
Gregory Ashton committed
138
        inputs = bilby_pipe.main.MainInput(args, self.unknown_args_list)
Gregory Ashton's avatar
Gregory Ashton committed
139
        self.assertEqual(inputs.detectors, ["H1", "L1"])
Gregory Ashton's avatar
Gregory Ashton committed
140

Gregory Ashton's avatar
Gregory Ashton committed
141
        args.detectors = ["H1", "l1"]
Gregory Ashton's avatar
Gregory Ashton committed
142 143
        inputs = bilby_pipe.main.MainInput(args, self.unknown_args_list)

144
        with self.assertRaises(BilbyPipeError):
Gregory Ashton's avatar
Gregory Ashton committed
145
            args.detectors = ["H1", "error"]
Gregory Ashton's avatar
Gregory Ashton committed
146 147
            inputs = bilby_pipe.main.MainInput(args, self.unknown_args_list)

148
    def test_create_summary_page(self):
Gregory Ashton's avatar
Gregory Ashton committed
149
        self.assertEqual(self.inputs.create_summary, self.args.create_summary)
150 151 152 153
        self.assertEqual(self.inputs.email, self.args.email)
        self.assertEqual(self.inputs.webdir, self.args.webdir)
        self.assertEqual(self.inputs.existing_dir, self.args.existing_dir)

Gregory Ashton's avatar
Gregory Ashton committed
154 155 156 157 158
    def test_n_simulation_setting(self):
        inputs = bilby_pipe.main.MainInput(self.args, self.unknown_args_list)
        inputs.n_simulation = 1
        self.assertEqual(inputs.n_simulation, 1)

159 160 161 162 163 164 165 166 167
        inputs.n_simulation = 3
        self.assertEqual(inputs.n_simulation, 3)

    # def test_n_simulation_None(self):
    #    args = self.args
    #    args.injection = True
    #    args.n_simulation = None
    #    inputs = bilby_pipe.main.MainInput(args, self.unknown_args_list)
    #    self.assertEqual(inputs.n_simulation, 0)
Gregory Ashton's avatar
Gregory Ashton committed
168 169 170 171 172

    def test_get_trigger_time_list(self):
        inputs = bilby_pipe.main.MainInput(self.args, self.unknown_args_list)

        inputs.trigger_time = 10
Avi Vajpeyi's avatar
Avi Vajpeyi committed
173 174 175 176 177 178
        self.assertEqual(
            bilby_pipe.job_creation.bilby_pipe_dag_creator.get_trigger_time_list(
                inputs
            ),
            [10],
        )
Gregory Ashton's avatar
Gregory Ashton committed
179 180

    def test_get_trigger_time_list_gps_file(self):
181
        inputs = bilby_pipe.main.MainInput(self.args, self.unknown_args_list)
Gregory Ashton's avatar
Gregory Ashton committed
182

183
        inputs.gps_file = self.test_gps_file
Avi Vajpeyi's avatar
Avi Vajpeyi committed
184
        A = bilby_pipe.job_creation.bilby_pipe_dag_creator.get_trigger_time_list(inputs)
Gregory Ashton's avatar
Gregory Ashton committed
185 186 187
        start_times = np.genfromtxt(self.test_gps_file)
        B = start_times + inputs.duration - inputs.post_trigger_duration
        self.assertTrue(np.all(A == B))
188

189 190 191 192 193 194 195 196 197 198 199 200 201
    def test_get_trigger_time_list_with_gaussian_noise_and_trigger_time(self):
        args = self.args
        args.trigger_time = 10
        args.gaussian_noise = True
        inputs = bilby_pipe.main.MainInput(self.args, self.unknown_args_list)

        self.assertEqual(
            bilby_pipe.job_creation.bilby_pipe_dag_creator.get_trigger_time_list(
                inputs
            ),
            [args.trigger_time],
        )

Gregory Ashton's avatar
Gregory Ashton committed
202 203 204 205 206 207
    def test_get_trigger_time_list_gaussian_noise(self):
        inputs = bilby_pipe.main.MainInput(self.args, self.unknown_args_list)

        inputs.gaussian_noise = True
        inputs.n_simulation = 3
        t = 0 + inputs.duration - inputs.post_trigger_duration
Avi Vajpeyi's avatar
Avi Vajpeyi committed
208 209 210 211 212 213
        self.assertTrue(
            bilby_pipe.job_creation.bilby_pipe_dag_creator.get_trigger_time_list(
                inputs
            ),
            [t] * 3,
        )
Gregory Ashton's avatar
Gregory Ashton committed
214 215

    def test_get_trigger_time_list_fail(self):
216 217
        inputs = bilby_pipe.main.MainInput(self.args, self.unknown_args_list)

Gregory Ashton's avatar
Gregory Ashton committed
218
        with self.assertRaises(BilbyPipeError):
Avi Vajpeyi's avatar
Avi Vajpeyi committed
219
            bilby_pipe.job_creation.bilby_pipe_dag_creator.get_trigger_time_list(inputs)
Gregory Ashton's avatar
Gregory Ashton committed
220

Gregory Ashton's avatar
Gregory Ashton committed
221 222
    def test_get_detectors_list(self):
        self.args.detectors = ["H1", "L1", "V1"]
Gregory Ashton's avatar
Gregory Ashton committed
223
        inputs = bilby_pipe.main.MainInput(self.args, self.unknown_args_list)
Avi Vajpeyi's avatar
Avi Vajpeyi committed
224 225 226
        det_list = bilby_pipe.job_creation.bilby_pipe_dag_creator.get_detectors_list(
            inputs
        )
Gregory Ashton's avatar
Gregory Ashton committed
227
        self.assertTrue(det_list, [["H1", "L1", "V1"]])
Gregory Ashton's avatar
Gregory Ashton committed
228

Gregory Ashton's avatar
Gregory Ashton committed
229 230
        self.args.detectors = ["H1", "L1", "V1"]
        self.args.coherence_test = True
Gregory Ashton's avatar
Gregory Ashton committed
231
        inputs = bilby_pipe.main.MainInput(self.args, self.unknown_args_list)
Avi Vajpeyi's avatar
Avi Vajpeyi committed
232 233 234
        det_list = bilby_pipe.job_creation.bilby_pipe_dag_creator.get_detectors_list(
            inputs
        )
Gregory Ashton's avatar
Gregory Ashton committed
235
        self.assertTrue(det_list, [["H1", "L1", "V1"], ["H1"], ["L1"], ["V1"]])
Gregory Ashton's avatar
Gregory Ashton committed
236

Gregory Ashton's avatar
Gregory Ashton committed
237

Gregory Ashton's avatar
Gregory Ashton committed
238
if __name__ == "__main__":
Gregory Ashton's avatar
Gregory Ashton committed
239
    unittest.main()