Coverage for pesummary/tests/injection_test.py: 100.0%

1# Licensed under an MIT style license -- see LICENSE.md 

2 

3import os 

4import shutil 

5import numpy as np 

6from .base import make_injection_file, testing_dir 

7from pesummary.gw.file.injection import GWInjection 

8import tempfile 

9 

10tmpdir = tempfile.TemporaryDirectory(prefix=".", dir=".").name 

11 

12__author__ = ["Charlie Hoy <charlie.hoy@ligo.org>"] 

13 

14 

15class TestInjection(object): 

16 """Class to test the Injection class for both the core and gw package 

17 """ 

18 def setup_method(self): 

19 """Setup the TestInjection class 

20 """ 

21 if not os.path.isdir(tmpdir): 

22 os.mkdir(tmpdir) 

23 

24 def teardown_method(self): 

25 """Remove the files and directories created from this class 

26 """ 

27 if os.path.isdir(tmpdir): 

28 shutil.rmtree(tmpdir) 

29 

30 def check(self, extension): 

31 """ 

32 """ 

33 if extension == "xml": 

34 ff = testing_dir + "/main_injection.xml" 

35 data = { 

36 'dec': [1.949725], 'geocent_time': [1186741861], 

37 'spin_2x': [0.0], 'spin_2y': [0.0], 'spin_2z': [0.0], 

38 'luminosity_distance': [139.7643], 'ra': [-1.261573], 

39 'spin_1y': [0.0], 'spin_1x': [0.0], 'spin_1z': [0.0], 

40 'psi': [1.75], 'phase': [0.0], 'iota': [1.0471976], 

41 'mass_1': [53.333332], 'mass_2': [26.666668], 

42 'symmetric_mass_ratio': [0.22222222], 'chirp_mass': [32.446098], 

43 'phase': [0.0] 

44 } 

45 else: 

46 ff, data = make_injection_file( 

47 extension=extension, return_filename=True, 

48 return_injection_dict=True, outdir=tmpdir 

49 ) 

50 inj = GWInjection.read(ff, conversion=False) 

51 assert all(param in data.keys() for param in inj.samples_dict.keys()) 

52 for param, value in inj.samples_dict.items(): 

53 np.testing.assert_almost_equal(value, data[param], 5) 

54 

55 def test_read_json(self): 

56 """Test that the Injection class can read in json formats 

57 """ 

58 self.check("json") 

59 

60 def test_read_dat(self): 

61 """Test that the Injection class can read in dat format 

62 """ 

63 self.check("dat") 

64 

65 def test_read_hdf5(self): 

66 """Test that the Injection class can read in hdf5 format 

67 """ 

68 self.check("hdf5") 

69 self.check("h5") 

70 

71 def test_read_xml(self): 

72 """Test that the Injection class can read in xml format 

73 """ 

74 self.check("xml")