Save beliefs data to file (#519)

Support showing beliefs in a custom resolution and/or timezone, and also saving shown beliefs data to a CSV file. * Allow saving shown beliefs to file Signed-off-by: F.N. Claessen <felix@seita.nl> * Support saving time series data in a specific resolution Signed-off-by: F.N. Claessen <felix@seita.nl> * Support saving time series data in a specific timezone, with the timezone of the first sensor as a default Signed-off-by: F.N. Claessen <felix@seita.nl> * black Signed-off-by: F.N. Claessen <felix@seita.nl> * CLI changelog entry Signed-off-by: F.N. Claessen <felix@seita.nl> * changelog entry Signed-off-by: F.N. Claessen <felix@seita.nl> * fix test Signed-off-by: F.N. Claessen <felix@seita.nl> * Clarify changelog entries Signed-off-by: F.N. Claessen <felix@seita.nl> Signed-off-by: F.N. Claessen <felix@seita.nl>
FlexMeasures · Nov 10, 2022 · e5a8a78 · e5a8a78
1 parent bfcdb98
commit e5a8a78
Show file tree

Hide file tree

Showing 4 changed files with 50 additions and 8 deletions.
diff --git a/documentation/changelog.rst b/documentation/changelog.rst
@@ -13,6 +13,7 @@ New features
 * Visually distinguish forecasts/schedules (dashed lines) from measurements (solid lines), and expand the tooltip with timing info regarding the forecast/schedule horizon or measurement lag [see `PR #503 <http://www.github.com/FlexMeasures/flexmeasures/pull/503>`_]
 * The asset page also allows to show sensor data from other assets that belong to the same account [see `PR #500 <http://www.github.com/FlexMeasures/flexmeasures/pull/500>`_]
 * Improved import of time series data from CSV file: 1) drop duplicate records with warning, and 2) allow configuring which column contains explicit recording times for each data point (use case: import forecasts) [see `PR #501 <http://www.github.com/FlexMeasures/flexmeasures/pull/501>`_]
+* The CLI command ``flexmeasures show beliefs`` supports showing beliefs data in a custom resolution and/or timezone, and also saving the shown beliefs data to a CSV file [see `PR #519 <http://www.github.com/FlexMeasures/flexmeasures/pull/519>`_]
 
 Bugfixes
 -----------

diff --git a/documentation/cli/change_log.rst b/documentation/cli/change_log.rst
@@ -4,9 +4,10 @@
 FlexMeasures CLI Changelog
 **********************
 
-since v0.12.0 | November xx, 2022
+since v0.12.0 | November XX, 2022
 =================================
 
+* Add ``--resolution``, ``--timezone`` and ``--to-file`` options to ``flexmeasures show beliefs``, to show beliefs data in a custom resolution and/or timezone, and also to save shown beliefs data to a CSV file.
 * Fix ``flexmeasures db-ops dump`` and ``flexmeasures db-ops restore`` incorrectly reporting a success when `pg_dump` and `pg_restore` are not installed.
 
 since v0.11.0 | August 28, 2022

diff --git a/flexmeasures/cli/data_show.py b/flexmeasures/cli/data_show.py
@@ -258,20 +258,45 @@ def list_data_sources():
     type=DataSourceIdField(),
     help="Source of the beliefs (an existing source id).",
 )
+@click.option(
+    "--resolution",
+    "resolution",
+    type=DurationField(),
+    required=False,
+    help="Resolution of the data. If not set, defaults to the minimum resolution of the sensor data.",
+)
+@click.option(
+    "--timezone",
+    "timezone",
+    type=str,
+    required=False,
+    help="Timezone of the data. If not set, defaults to the timezone of the first non-empty sensor.",
+)
+@click.option(
+    "--to-file",
+    "filepath",
+    required=False,
+    type=str,
+    help="Set a filepath to store the beliefs as a CSV file.",
+)
 def plot_beliefs(
     sensors: List[Sensor],
     start: datetime,
     duration: timedelta,
+    resolution: Optional[timedelta],
+    timezone: Optional[str],
     belief_time_before: Optional[datetime],
     source: Optional[DataSource],
+    filepath: Optional[str],
 ):
     """
-    Show a simple plot of belief data directly in the terminal.
+    Show a simple plot of belief data directly in the terminal, and optionally, save the data to a CSV file.
     """
     sensors = list(sensors)
-    minimum_resampling_resolution = determine_minimum_resampling_resolution(
-        [sensor.event_resolution for sensor in sensors]
-    )
+    if resolution is None:
+        resolution = determine_minimum_resampling_resolution(
+            [sensor.event_resolution for sensor in sensors]
+        )
 
     # query data
     beliefs_by_sensor = TimedBelief.search(
@@ -281,7 +306,7 @@ def plot_beliefs(
         beliefs_before=belief_time_before,
         source=source,
         one_deterministic_belief_per_event=True,
-        resolution=minimum_resampling_resolution,
+        resolution=resolution,
         sum_multiple=False,
     )
     # only keep non-empty
@@ -303,6 +328,11 @@ def plot_beliefs(
     df = pd.concat([simplify_index(df) for df in beliefs_by_sensor.values()], axis=1)
     df.columns = beliefs_by_sensor.keys()
 
+    # Convert to the requested or default timezone
+    if timezone is not None:
+        timezone = sensors[0].timezone
+    df.index = df.index.tz_convert(timezone)
+
     # Build title
     if len(sensors) == 1:
         title = f"Beliefs for Sensor '{sensors[0].name}' (Id {sensors[0].id}).\n"
@@ -313,7 +343,7 @@ def plot_beliefs(
         title += f"\nOnly beliefs made before: {belief_time_before}."
     if source:
         title += f"\nSource: {source.description}"
-    title += f"\nThe time resolution (x-axis) is {naturaldelta(minimum_resampling_resolution)}."
+    title += f"\nThe time resolution (x-axis) is {naturaldelta(resolution)}."
 
     uniplot.plot(
         [df[col] for col in df.columns],
@@ -325,6 +355,12 @@ def plot_beliefs(
         if shared_unit
         else [s.name + f" (in {s.unit})" for s in sensors],
     )
+    if filepath is not None:
+        df.columns = pd.MultiIndex.from_arrays(
+            [df.columns, [df.sensor.unit for df in beliefs_by_sensor.values()]]
+        )
+        df.to_csv(filepath)
+        click.echo("Data saved to file.")
 
 
 app.cli.add_command(fm_show_data)
diff --git a/flexmeasures/cli/tests/test_data_show.py b/flexmeasures/cli/tests/test_data_show.py
@@ -1,4 +1,5 @@
 import pytest
+from click.formatting import wrap_text
 
 from flexmeasures.data.models.time_series import Sensor
 from flexmeasures.cli.tests.utils import get_click_commands
@@ -121,5 +122,8 @@ def test_cli_help(app):
     for cmd in get_click_commands(data_show):
         result = runner.invoke(cmd, ["--help"])
         assert "Usage" in result.output
-        assert cmd.__doc__.strip() in result.output
+        assert (
+            wrap_text(cmd.__doc__.strip(), initial_indent="  ", subsequent_indent="  ")
+            in result.output
+        )
         assert result.exit_code == 0