Coverage for pesummary/gw/file/formats/princeton.py: 86.7%

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

2 

3from pesummary.core.file.formats.numpy import load 

4 

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

6# column names taken from https://github.com/jroulet/O2_samples 

7_O1O2_column_names = [ 

8 'chirp_mass', 'symmetric_mass_ratio', 'spin_1z', 'spin_2z', 'ra', 'dec', 

9 'psi', 'iota', 'phase', 'geocent_time', 'luminosity_distance' 

10] 

11_column_names = { 

12 "O1O2": _O1O2_column_names 

13} 

14 

15 

16def read_princeton(path, type="O1O2", column_names=None, **kwargs): 

17 """Grab the parameters and samples in a princeton file 

18 

19 Parameters 

20 ---------- 

21 path: str 

22 path to the result file you wish to read in 

23 type: str, optional 

24 the type of file being loaded. This affects how the columns are 

25 assigned. Default O1O2 

26 column_names: list, optional 

27 list of column names corresponding to the loaded numpy array 

28 """ 

29 if column_names is None and type not in _column_names: 

30 raise ValueError( 

31 "Please specify the type of file you are loading. This is used to " 

32 "make sure that the columns are correctly assigned. The list of " 

33 "available types are {}".format(", ".join(_column_names.keys())) 

34 ) 

35 elif column_names is None: 

36 column_names = _column_names[type] 

37 return _read_princeton(path, column_names=column_names, **kwargs) 

38 

39 

40def _read_princeton(path, column_names, **kwargs): 

41 """Grab the samples from a princeton file 

42 

43 Parameters 

44 ---------- 

45 path: str 

46 path to the result file you wish to read in 

47 column_names: list 

48 list of column names corresponding to the loaded numpy array 

49 """ 

50 samples = load(path) 

51 if len(column_names) != len(samples[0]): 

52 raise ValueError( 

53 "The number of column names does not match the number of columns " 

54 "in the loaded numpy array. Did you pick the incorrect 'type'?" 

55 ) 

56 return column_names, samples