Coverage for pesummary/cli/summarydetchar.py: 88.9%

1#! /usr/bin/env python 

2 

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

4 

5import os 

6from pesummary.gw.file.read import read 

7from pesummary.gw.plots import detchar 

8from pesummary.utils.exceptions import InputError 

9from pesummary.utils.utils import make_dir 

10from pesummary.gw.cli.parser import ArgumentParser as _ArgumentParser 

11 

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

13__doc__ = """This executable is used to generate plots associated with the 

14detectors""" 

15 

16 

17class ArgumentParser(_ArgumentParser): 

18 def _pesummary_options(self): 

19 options = super(ArgumentParser, self)._pesummary_options() 

20 options.update( 

21 { 

22 "--plot": { 

23 "choices": ["spectrogram", "omegascan"], 

24 "default": "2d_contour", 

25 "help": "name of the publication plot you wish to produce", 

26 }, 

27 "--gps": { 

28 "help": "GPS time to centre the omegascan around", 

29 }, 

30 "--vmin": { 

31 "default": 0, 

32 "help": "minimum for the omegascan colormap" 

33 }, 

34 "--vmax": { 

35 "default": 25, 

36 "help": "maximum for the omegascan colormap" 

37 }, 

38 "--window": { 

39 "default": 4, 

40 "help": "window around gps time to generate omegascan for", 

41 } 

42 } 

43 ) 

44 return options 

45 

46 

47def get_maxL_time(samples): 

48 """Return the maxL time stored in the samples 

49 

50 Parameters 

51 ---------- 

52 samples: str 

53 path to a samples file 

54 """ 

55 f = read(samples) 

56 samples_dict = f.samples_dict 

57 return samples_dict["geocent_time"].maxL 

58 

59 

60def read_strain(dictionary): 

61 """Read the gwdata strain and return a gwpy.timeseries.TimeSeries object 

62 

63 Parameters 

64 ---------- 

65 dictionary: dict 

66 dictionary of channels and cache files 

67 """ 

68 from pesummary.gw.file.strain import StrainDataDict 

69 

70 for i in dictionary.keys(): 

71 if not os.path.isfile(dictionary[i]): 

72 raise InputError( 

73 "The file {} does not exist. Please check the path to " 

74 "your strain file".format(dictionary[i]) 

75 ) 

76 timeseries = StrainDataDict.read(dictionary) 

77 return timeseries 

78 

79 

80def make_spectrogram_plot(opts): 

81 """Make a spectrogram plot 

82 """ 

83 gwdata = read_strain(opts.gwdata) 

84 figs = detchar.spectrogram(gwdata) 

85 for det, fig in figs.items(): 

86 fig.savefig( 

87 os.path.join( 

88 opts.webdir, "spectrogram_{}.png".format(det) 

89 ) 

90 ) 

91 fig.close() 

92 

93 

94def make_omegascan_plot(opts): 

95 """Make an omegascan plot. If gps is None, centre around maxL from samples 

96 """ 

97 if opts.gps is None: 

98 opts.gps = get_maxL_time(opts.samples[0]) 

99 gwdata = read_strain(opts.gwdata) 

100 figs = detchar.omegascan( 

101 gwdata, float(opts.gps), window=float(opts.window), 

102 vmin=float(opts.vmin), vmax=float(opts.vmax) 

103 ) 

104 for det, fig in figs.items(): 

105 fig.savefig( 

106 os.path.join( 

107 opts.webdir, "omegascan_{}.png".format(det) 

108 ) 

109 ) 

110 fig.close() 

111 

112 

113def main(args=None): 

114 """Top level interface for `summarydetchar` 

115 """ 

116 parser = ArgumentParser(description=__doc__) 

117 parser.add_known_options_to_parser( 

118 [ 

119 "--webdir", "--samples", "--gwdata", "--plot", "--gps", "--vmin", 

120 "--vmax", "--window" 

121 ] 

122 ) 

123 opts = parser.parse_args(args=args) 

124 make_dir(opts.webdir) 

125 func_map = {"spectrogram": make_spectrogram_plot, 

126 "omegascan": make_omegascan_plot} 

127 func_map[opts.plot](opts)