Coverage for pesummary/core/notebook/notebook.py: 84.7%
59 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-05-02 08:42 +0000
1# Licensed under an MIT style license -- see LICENSE.md
3try:
4 import nbformat as nbf
5except ImportError:
6 raise ImportError("'nbformat' is required for this module")
7import os
9__author__ = ["Charlie Hoy <charlie.hoy@ligo.org>"]
12class NoteBook(object):
13 """Class to handle the creation of a jupyter notebook
15 Attributes
16 ----------
17 cells: dict
18 dictionary containing the cell contents keyed by their cell number
19 """
20 def __init__(self):
21 self.cells = {}
23 def add_cell(self, string, markdown=False, code=False):
24 """Add cell to the cell dictionary
26 Parameters
27 ----------
28 string: str
29 string containing the cell contents
30 markdown: Bool, optional
31 if True, treat the cell as a markdown cell. Default False
32 code: Bool, optional
33 if True, treat the cell as a code cell. Default False
34 """
35 cell_number = len(self.cells)
36 if markdown and not code:
37 self.cells[cell_number] = nbf.v4.new_markdown_cell(string)
38 elif code and not markdown:
39 self.cells[cell_number] = nbf.v4.new_code_cell(string)
40 else:
41 raise ValueError("String must either be markdown or code not both")
43 def write(self, filename="posterior_samples.ipynb", outdir="./"):
44 """Save the cells to file
46 Parameters
47 ----------
48 cells: dict
49 dictionary containing the cell information. Keyed by their cell
50 number
51 filename: str, optional
52 name of the file you wish to save the notebook to. Default
53 posterior_samples.ipynb
54 outdir: str, optional
55 the directory to save the file. Default ./
56 """
57 nb = nbf.v4.new_notebook()
58 nb['cells'] = [self.cells[num] for num in range(len(self.cells))]
59 nbf.write(nb, os.path.join(outdir, filename))
62def imports(
63 magic="%matplotlib inline", comment=None,
64 module_imports=["pesummary", "pesummary.io:read"],
65 text="First we import the key python modules", extra_lines=[]
66):
67 """Return a string containing the key imports
69 Parameters
70 ----------
71 comment: str, optional
72 comment to add at the top of the string. Default None
73 magic: str, optional
74 magic statement to include in the string. Default '%matplotlib inline'
75 module_imports: list, optional
76 list of modules you wish to import. If a colon is in the import string
77 this is interpreted as, 'from %s import %s' % tuple(import.split(':'))'
78 text: str, optional
79 Markdown text explaining the import cell below. Default
80 'First we import the key python modules'
81 extra_lines: list, optional
82 optional lines to add to the end of the string
83 """
84 string = ""
85 if comment is not None:
86 string += comment + "\n"
87 if magic is not None:
88 string += magic + "\n"
89 for _import in module_imports:
90 if ":" not in _import:
91 string += "import {}\n".format(_import)
92 else:
93 string += "from %s import %s\n" % tuple(_import.split(":"))
94 string += "\n".join(extra_lines)
95 if text is not None:
96 return [text, string]
97 return [string]
100def pesummary_read(
101 path, text="We now load in the file using the 'pesummary' read function",
102 read_variable="data"
103):
104 """Return a string containing the function to read a file
106 Parameters
107 ----------
108 path: str
109 path to the result file you wish to load
110 text: str, optional
111 Markdown text explaining the import cell below. Default
112 'We now load in the file using the 'pesummary' read function'
113 read_variable: str, optional
114 name of the variable to assign to the read result file. Default 'data'
115 """
116 string = "file_name = '{}'\n{} = read(file_name)".format(path, read_variable)
117 if text is not None:
118 return [text, string]
119 return [string]
122def posterior_samples(
123 read_variable, metafile=False, default_analysis=None,
124 print_parameters=True, samples_variable="posterior_samples", text=(
125 "The posterior samples can be extracted through the `samples_dict` "
126 "property. These posterior samples are stored in a custom "
127 "table structure"
128 )
129):
130 """Return a string showing how to extract the posterior samples from a
131 result file
133 Parameters
134 ----------
135 read_variable: str
136 name of the read object
137 metafile: Bool, optional
138 if True, the result file is a pesummary meta file
139 default_analysis: str, optional
140 default analysis to use when `metafile=True`
141 print_parameters: Bool, optional
142 if True, print the parameters stored in the table
143 samples_variable: str, optional
144 name of the SamplesDict class
145 text: str, optional
146 Markdown text explaining how to extract posterior samples from the file
147 """
148 string = "samples_dict = {}.samples_dict\n".format(read_variable)
149 if metafile:
150 if default_analysis is None:
151 raise ValueError("Please provide a default analysis")
152 string += "{} = samples_dict['{}']\n".format(
153 samples_variable, default_analysis
154 )
155 if print_parameters:
156 string += "parameters = list({}.keys())\n".format(samples_variable)
157 string += "print(parameters)"
158 if text is not None:
159 return [text, string]
160 return [string]
163def samples_dict_plot(
164 samples_variable, plot_args=[], plot_kwargs={}, extra_lines=[],
165 text="As an example, we now plot the posterior samples"
166):
167 """Return a string containing the function to generate a plot
169 Parameters
170 ----------
171 samples_variable: str
172 name of the SamplesDict class
173 plot_args: list, optional
174 arguments for the `.plot()` method
175 plot_kwargs: dict, optional
176 kwargs for the `.plot()` method.
177 extra_lines: list, optional
178 additional lines to add to the end of the string
179 text: str, optional
180 Markdown text explaining the plot
181 """
182 args = ", ".join(plot_args)
183 kwargs = ", ".join([f"{key}={item}" for key, item in plot_kwargs.items()])
184 string = "fig = {}.plot({})\n".format(
185 samples_variable, "%s, %s" % (args, kwargs) if len(args) else kwargs
186 )
187 string += "\n".join(extra_lines)
188 if text is not None:
189 return [text, string]
190 return [string]