Skip to content

Commit 09f63ea

Browse files
SBlokhuizenSBlokhuizen
committed
update docstrings
1 parent cd1e6b7 commit 09f63ea

File tree

3 files changed

+158
-31
lines changed

3 files changed

+158
-31
lines changed

‎pyproject.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ dependencies = [
2222
"pygments",
2323
"plotly",
2424
"holoviews",
25-
"kaleido",
25+
"kaleido", # required for exporting PNG to disk
2626
]
2727
dynamic = ["version"]
2828

‎waveform_editor/cli.py‎

Lines changed: 84 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,19 @@ def print_version():
6262
@click.option("--times", type=click.Path(exists=True))
6363
@click.option("--num_interp", type=int)
6464
def export_csv(yaml, output, times, num_interp):
65-
"""Export waveform data to a CSV file."""
65+
"""Export waveform data to a CSV file.
66+
67+
\b
68+
Arguments:
69+
yaml: Path to the waveform YAML file.
70+
output: Path where the CSV file will be saved.
71+
72+
\b
73+
Options:
74+
--times: CSV file containing a custom time array (column-based)
75+
--num_interp: Number of points for linear interpolation (only used if --times
76+
is not provided).
77+
"""
6678
exporter = setup_exporter(yaml, times, num_interp)
6779
exporter.to_csv(output)
6880

@@ -73,7 +85,19 @@ def export_csv(yaml, output, times, num_interp):
7385
@click.option("--times", type=click.Path(exists=True))
7486
@click.option("--num_interp", type=int)
7587
def export_png(yaml, output, times, num_interp):
76-
"""Export waveform data to a CSV file."""
88+
"""Export waveform data to a PNG file.
89+
90+
\b
91+
Arguments:
92+
yaml: Path to the waveform YAML file.
93+
output: Path where the PNG file will be saved.
94+
95+
\b
96+
Options:
97+
--times: CSV file containing a custom time array (column-based).
98+
--num_interp: Number of points for linear interpolation (only used if --times
99+
is not provided).
100+
"""
77101
exporter = setup_exporter(yaml, times, num_interp)
78102
exporter.to_png(output)
79103

@@ -85,21 +109,69 @@ def export_png(yaml, output, times, num_interp):
85109
@click.option("--times", type=click.Path(exists=True))
86110
@click.option("--num_interp", type=int)
87111
def export_ids(yaml, uri, dd_version, times, num_interp):
88-
"""Export waveform data to an IDS."""
112+
"""Export waveform data to an IDS.
113+
114+
\b
115+
Arguments:
116+
yaml: Path to the waveform YAML file.
117+
uri: URI containing the IDS, and path to export to. (See below for examples)
118+
119+
\b
120+
Options:
121+
--dd-version: Data Dictionary version to use for the IDS export, if not provided
122+
IMASPy's default DD-version will be used.
123+
--times: CSV file containing a custom time array (column-based).
124+
--num_interp: Number of points for linear interpolation (only used if --times
125+
is not provided).
126+
127+
\b
128+
Example URIs:
129+
- imas:hdf5?path=./testdb#ec_launchers/beam(1)/power_launched
130+
- imas:hdf5?path=./testdb#ec_launchers:1/beam(1)/power_launched
131+
- imas:hdf5?path=./testdb#equilibrium/time_slice()/boundary/elongation
132+
"""
89133
exporter = setup_exporter(yaml, times, num_interp)
90134
exporter.to_ids(uri, dd_version=dd_version)
91135

92136

93137
def setup_exporter(yaml, times, num_interp):
138+
"""Initialize and return a WaveformExporter.
139+
140+
Args:
141+
yaml: Path to the waveform YAML file.
142+
times: Path to a CSV file containing a custom time array.
143+
num_interp: Number of points for linear interpolation (only used if `times`
144+
is None).
145+
Returns:
146+
An instance of the WaveformExporter configured with the waveform.
147+
"""
148+
94149
waveform = load_waveform_from_yaml(yaml)
95150
time_array = load_time_array(times, waveform, num_interp)
96151
exporter = WaveformExporter(waveform, times=time_array)
97152
return exporter
98153

99154

100155
def load_time_array(times, waveform, num_interp):
101-
"""Load time array from CSV file or use default linear interpolation."""
156+
"""Load time array from CSV file or use default linear interpolation.
157+
158+
Arguments:
159+
times: Path to a CSV file containing a custom time array, or None to use linear
160+
interpolation.
161+
waveform: Waveform to load.
162+
num_interp: Number of points for linear interpolation (only used if `times`
163+
is None).
164+
165+
Returns:
166+
A numpy array containing the time values.
167+
"""
102168
if times:
169+
if num_interp:
170+
click.secho(
171+
"Both `--num_interp` and `--times` were set. The provided times will "
172+
"be used, and `num_interp` will be ignored.",
173+
fg="yellow",
174+
)
103175
try:
104176
# assuming single column format
105177
with open(times, newline="") as csvfile:
@@ -124,6 +196,14 @@ def load_time_array(times, waveform, num_interp):
124196

125197

126198
def load_waveform_from_yaml(yaml_file):
199+
"""Load a waveform object from a YAML file.
200+
201+
Arguments:
202+
yaml_file: Path to the YAML file.
203+
204+
Returns:
205+
The waveform parsed from the YAML file.
206+
"""
127207
with open(yaml_file) as file:
128208
yaml_str = file.read()
129209
yaml_parser = YamlParser()

‎waveform_editor/waveform_exporter.py‎

Lines changed: 73 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,18 @@ def __init__(self, waveform, times):
1919
self.value_label = "Value [unit]"
2020

2121
def parse_uri(self, uri):
22+
"""Parse URI into its constituents.
23+
24+
Args:
25+
uri: String containing the URI.
26+
27+
Returns:
28+
Constituents of the URI:
29+
- The scheme, backend and query parts of the URI
30+
- The name of the IDS.
31+
- The occurrence number (defaults to 0 if not provided)
32+
- Path of the IDS quantity to export to
33+
"""
2234
uri_entry, fragment = uri.split("#")
2335
fragment_parts = fragment.split("/")
2436
idsname_part = fragment_parts[0]
@@ -35,6 +47,13 @@ def parse_uri(self, uri):
3547
return uri_entry, ids_name, occurrence, ids_path
3648

3749
def to_ids(self, uri, dd_version=None):
50+
"""Export the waveform to an IDS.
51+
52+
Args:
53+
uri: URI containing scheme, backend, query and fragment parts.
54+
dd_version: The data dictionary version to export to. If None, IMASPy's
55+
default version will be used.
56+
"""
3857
uri_entry, uri_ids, occurrence, path = self.parse_uri(uri)
3958
entry = imaspy.DBEntry(uri_entry, "r", dd_version=dd_version)
4059
ids = entry.get(uri_ids, occurrence, autoconvert=False)
@@ -63,39 +82,23 @@ def to_ids(self, uri, dd_version=None):
6382
entry.put(ids)
6483
entry.close()
6584

66-
def _fill_flt_0d(self, ids, path, is_homogeneous):
67-
aos_path, remaining_path = path.split("()")
68-
aos_path = aos_path.strip("/")
69-
remaining_path = remaining_path.strip("/")
70-
aos = ids[aos_path]
71-
aos.resize(len(self.times))
72-
73-
for i, time in enumerate(self.times):
74-
if aos[i][remaining_path].data_type == "FLT_0D":
75-
aos[i][remaining_path] = self.values[i]
76-
if not is_homogeneous:
77-
aos[i].time = time
78-
else:
79-
raise NotImplementedError("Should be float 0d")
80-
81-
def _fill_flt_1d(self, ids, path, is_homogeneous):
82-
quantity = ids[path]
83-
struct_ref = quantity.metadata.structure_reference
84-
if struct_ref == "signal_flt_1d":
85-
quantity.data = self.values
86-
if not is_homogeneous:
87-
quantity.time = self.times
88-
else:
89-
raise NotImplementedError("Invalid data")
90-
9185
def to_csv(self, file_path):
86+
"""Export the waveform to a CSV.
87+
88+
Args:
89+
file_path: The file path and name to store the CSV to.
90+
"""
9291
with open(file_path, mode="w", newline="") as file:
9392
writer = csv.writer(file)
9493
writer.writerow([self.time_label, self.value_label])
9594
writer.writerows(zip(self.times, self.values))
9695

9796
def to_png(self, file_path):
98-
"""Export waveform data as a PNG plot using Plotly Express."""
97+
"""Export the waveform to a PNG.
98+
99+
Args:
100+
file_path: The file path and name to store the PNG to.
101+
"""
99102

100103
fig = go.Figure(
101104
data=go.Scatter(
@@ -113,3 +116,47 @@ def to_png(self, file_path):
113116
# TODO: implement export to XML format
114117
# def to_pcssp(self, file_path)
115118
# pass
119+
120+
def _fill_flt_0d(self, ids, path, is_homogeneous):
121+
"""Fill a FLT_0D IDS quantity in an IDS.
122+
123+
It is assumed that the time dependent AoS is provided in the path using `()`,
124+
for example:
125+
126+
imas:hdf5?path=./test_equilibrium#equilibrium/time_slice()/boundary/elongation
127+
128+
Arguments:
129+
ids: The IDS to fill.
130+
path: The path to the FLT_0D quantity to fill.
131+
is_homogeneous: Whether to fill the local time array, or the ids.time array
132+
"""
133+
aos_path, remaining_path = path.split("()")
134+
aos_path = aos_path.strip("/")
135+
remaining_path = remaining_path.strip("/")
136+
aos = ids[aos_path]
137+
aos.resize(len(self.times))
138+
139+
for i, time in enumerate(self.times):
140+
if aos[i][remaining_path].data_type == "FLT_0D":
141+
aos[i][remaining_path] = self.values[i]
142+
if not is_homogeneous:
143+
aos[i].time = time
144+
else:
145+
raise NotImplementedError("Should be float 0d")
146+
147+
def _fill_flt_1d(self, ids, path, is_homogeneous):
148+
"""Fill a FLT_1D IDS quantity in an IDS.
149+
150+
Arguments:
151+
ids: The IDS to fill.
152+
path: The path to the FLT_1D quantity to fill.
153+
is_homogeneous: Whether to fill the local time array, or the ids.time array
154+
"""
155+
quantity = ids[path]
156+
struct_ref = quantity.metadata.structure_reference
157+
if struct_ref == "signal_flt_1d":
158+
quantity.data = self.values
159+
if not is_homogeneous:
160+
quantity.time = self.times
161+
else:
162+
raise NotImplementedError("Invalid data")

0 commit comments

Comments
 (0)