plots: Revise plotting function interfaces to return JSON string easily

examples/development/entrypoint_plotting.py

@@ -1,6 +1,7 @@
 """Plotting entrypoint development script"""
 
 import mesmo
+from mesmo import plots
 
 
 def main():
@@ -10,28 +11,31 @@ def main():
 
     mesmo.utils.cleanup()
     results_path = mesmo.utils.get_results_path("run_operation_problem", scenario_name)
-    results = mesmo.api.run_nominal_operation_problem(
+    results_raw = mesmo.api.run_nominal_operation_problem(
         scenario_name, results_path=results_path, store_results=False, recreate_database=False
     )
-    run_results = results.get_run_results()
+    results = results_raw.get_run_results()
 
     # Roundtrip save/load to/from JSON, just for demonstration
-    with open(results_path / "run_results.json", "w", encoding="utf-8") as file:
+    with open(results_path / "results.json", "w", encoding="utf-8") as file:
         print("Dumping results to file")
-        file.write(run_results.model_dump_json())
-    with open(results_path / "run_results.json", "r", encoding="utf-8") as file:
+        file.write(results.model_dump_json())
+    with open(results_path / "results.json", "r", encoding="utf-8") as file:
         print("Loading results from file")
-        run_results = mesmo.data_models.RunResults.model_validate_json(file.read())
-
-    # TODO: Return JSON object, function should take run_id as input
-    mesmo.plots.der_active_power_time_series(run_results, results_path)
-    mesmo.plots.der_reactive_power_time_series(run_results, results_path)
-    mesmo.plots.der_apparent_power_time_series(run_results, results_path)
-    mesmo.plots.der_aggregated_active_power_time_series(run_results, results_path)
-    mesmo.plots.der_aggregated_reactive_power_time_series(run_results, results_path)
-    mesmo.plots.der_aggregated_apparent_power_time_series(run_results, results_path)
-    mesmo.plots.node_voltage_per_unit_time_series(run_results, results_path)
-    mesmo.plots.node_aggregated_voltage_per_unit_time_series(run_results, results_path)
+        results = mesmo.data_models.RunResults.model_validate_json(file.read())
+
+    # Sample plotting to file, just for demonstration
+    plots.plot_to_file(plots.der_active_power_time_series, results=results, results_path=results_path)
+    plots.plot_to_file(plots.der_reactive_power_time_series, results=results, results_path=results_path)
+    plots.plot_to_file(plots.der_apparent_power_time_series, results=results, results_path=results_path)
+    plots.plot_to_file(plots.der_aggregated_active_power_time_series, results=results, results_path=results_path)
+    plots.plot_to_file(plots.der_aggregated_reactive_power_time_series, results=results, results_path=results_path)
+    plots.plot_to_file(plots.der_aggregated_apparent_power_time_series, results=results, results_path=results_path)
+    plots.plot_to_file(plots.node_voltage_per_unit_time_series, results=results, results_path=results_path)
+    plots.plot_to_file(plots.node_aggregated_voltage_per_unit_time_series, results=results, results_path=results_path)
+
+    # Sample JSON return
+    print(plots.plot_to_json(plots.der_active_power_time_series, results=results))
 
     # Print results path.
     mesmo.utils.launch(results_path)

mesmo/data_models/results.py

@@ -1,5 +1,5 @@
 """Results data models."""
-from typing import Annotated, Optional
+from typing import Optional
 
 import pandas as pd
 

mesmo/plots/__init__.py

@@ -1,12 +1,13 @@
 """Plotting function collection."""
 
+from .plots import plot_to_figure, plot_to_file, plot_to_json
 from .time_series import (
     der_active_power_time_series,
-    der_reactive_power_time_series,
-    der_apparent_power_time_series,
     der_aggregated_active_power_time_series,
-    der_aggregated_reactive_power_time_series,
     der_aggregated_apparent_power_time_series,
-    node_voltage_per_unit_time_series,
+    der_aggregated_reactive_power_time_series,
+    der_apparent_power_time_series,
+    der_reactive_power_time_series,
     node_aggregated_voltage_per_unit_time_series,
+    node_voltage_per_unit_time_series,
 )

mesmo/plots/plot_utils.py

@@ -2,14 +2,29 @@
 
 import json
 import pathlib
+
 import plotly.graph_objects as go
 import plotly.io as pio
 
 import mesmo.config
 import mesmo.utils
 
 
-def write_figure_plotly(
+def get_plotly_figure_json(figure: go.Figure) -> str:
+    """Get JSON string representation of plotly figure.
+
+    Args:
+        figure (go.Figure): Figure for which the JSON representation is generated
+
+    Returns:
+        str: JSON representation of given figure
+    """
+    json_dict = json.loads(pio.to_json(figure))
+    json_dict["layout"].pop("template")  # Exclude template information to minify JSON
+    return json.dumps(json_dict)
+
+
+def write_plotly_figure_file(
     figure: go.Figure,
     results_path: pathlib.Path,
     file_format=mesmo.config.config["plots"]["file_format"],
@@ -47,4 +62,4 @@ def write_figure_plotly(
 
     # Additionally to the requested format, also output at plottable item in JSON format.
     if file_format != "json":
-        write_figure_plotly(figure, results_path, "json", width, height)
+        write_plotly_figure_file(figure, results_path, "json", width, height)

mesmo/plots/plots.py

@@ -0,0 +1,62 @@
+"""Interfaces to plotting functions."""
+
+import pathlib
+from typing import Callable
+
+import plotly.graph_objects as go
+
+from mesmo import data_models
+from mesmo.plots import plot_utils
+
+
+def plot_to_figure(
+    *plot_functions: Callable[[go.Figure, data_models.RunResults], go.Figure], results: data_models.RunResults
+) -> go.Figure:
+    """Generate new plotly figure and apply given plotting function(s) for given run results.
+
+    Args:
+        Callable: Plotting function
+        data_models.RunResults: MESMO run results as input for the plotting function
+
+    Returns:
+        go.Figure: Plotly figure containing the generated plot
+    """
+    figure = go.Figure()
+    for plot_function in plot_functions:
+        figure = plot_function(figure, results)
+    return figure
+
+
+def plot_to_json(
+    *plot_functions: Callable[[go.Figure, data_models.RunResults], go.Figure], results: data_models.RunResults
+) -> str:
+    """Generate new plotly figure and apply given plotting function(s) for given run results. Output the final figure
+    as JSON string.
+
+    Args:
+        Callable: Plotting function
+        data_models.RunResults: MESMO run results as input for the plotting function
+
+    Returns:
+        str: JSON string containing the generated plot
+    """
+    return plot_utils.get_plotly_figure_json(plot_to_figure(*plot_functions, results=results))
+
+
+def plot_to_file(
+    *plot_functions: Callable[[go.Figure, data_models.RunResults], go.Figure],
+    results: data_models.RunResults,
+    results_path: pathlib.Path,
+):
+    """Generate new plotly figure and apply given plotting function(s) for given run results. Out put the final figure
+    to file.
+
+    Args:
+        Callable: Plotting function
+        data_models.RunResults: MESMO run results as input for the plotting function
+        pathlib.Path: Results file output path
+    """
+    filename = plot_functions[0].__name__
+    plot_utils.write_plotly_figure_file(
+        plot_to_figure(*plot_functions, results=results), results_path=results_path / filename
+    )

mesmo/plots/time_series.py

@@ -1,21 +1,18 @@
 """Timeseries-base plotting functions."""
 
 import numpy as np
-import pathlib
 import plotly.graph_objects as go
 
 from mesmo import data_models
-from mesmo.plots import plot_utils, constants
+from mesmo.plots import constants
 
 
-def der_active_power_time_series(results: data_models.RunResults, results_path: pathlib.Path):
+def der_active_power_time_series(figure: go.Figure, results: data_models.RunResults) -> go.Figure:
     title = f"{constants.ValueLabels.ACTIVE_POWER} per DER"
-    filename = der_active_power_time_series.__name__
     x_label = constants.ValueLabels.TIME
     y_label = f"{constants.ValueLabels.ACTIVE_POWER} [{constants.ValueUnitLabels.WATT}]"
     legend_title = constants.ValueLabels.DERS
 
-    figure = go.Figure()
     for der_type, der_name in results.der_model_set_index.ders:
         values = results.der_operation_results.der_active_power_vector.loc[:, (der_type, der_name)]
         figure.add_trace(go.Scatter(x=values.index, y=values.values, name=f"{der_name} ({der_type})"))
@@ -25,17 +22,15 @@ def der_active_power_time_series(results: data_models.RunResults, results_path:
         yaxis_title=y_label,
         legend=go.layout.Legend(title=legend_title, x=0.99, xanchor="auto", y=0.99, yanchor="auto"),
     )
-    plot_utils.write_figure_plotly(figure, results_path / filename)
+    return figure
 
 
-def der_reactive_power_time_series(results: data_models.RunResults, results_path: pathlib.Path):
+def der_reactive_power_time_series(figure: go.Figure, results: data_models.RunResults) -> go.Figure:
     title = f"{constants.ValueLabels.REACTIVE_POWER} per DER"
-    filename = der_reactive_power_time_series.__name__
     x_label = constants.ValueLabels.TIME
     y_label = f"{constants.ValueLabels.REACTIVE_POWER} [{constants.ValueUnitLabels.VOLT_AMPERE_REACTIVE}]"
     legend_title = constants.ValueLabels.DERS
 
-    figure = go.Figure()
     for der_type, der_name in results.der_model_set_index.ders:
         values = results.der_operation_results.der_reactive_power_vector.loc[:, (der_type, der_name)]
         figure.add_trace(go.Scatter(x=values.index, y=values.values, name=f"{der_name} ({der_type})"))
@@ -45,17 +40,16 @@ def der_reactive_power_time_series(results: data_models.RunResults, results_path
         yaxis_title=y_label,
         legend=go.layout.Legend(title=legend_title, x=0.99, xanchor="auto", y=0.99, yanchor="auto"),
     )
-    plot_utils.write_figure_plotly(figure, results_path / filename)
+    return figure
 
 
-def der_apparent_power_time_series(results: data_models.RunResults, results_path: pathlib.Path):
+def der_apparent_power_time_series(figure: go.Figure, results: data_models.RunResults) -> go.Figure:
     title = f"{constants.ValueLabels.APPARENT_POWER} per DER"
     filename = der_apparent_power_time_series.__name__
     x_label = constants.ValueLabels.TIME
     y_label = f"{constants.ValueLabels.APPARENT_POWER} [{constants.ValueUnitLabels.VOLT_AMPERE}]"
     legend_title = constants.ValueLabels.DERS
 
-    figure = go.Figure()
     for der_type, der_name in results.der_model_set_index.ders:
         # TODO: Add apparent power in result directly
         values = np.sqrt(
@@ -69,17 +63,15 @@ def der_apparent_power_time_series(results: data_models.RunResults, results_path
         yaxis_title=y_label,
         legend=go.layout.Legend(title=legend_title, x=0.99, xanchor="auto", y=0.99, yanchor="auto"),
     )
-    plot_utils.write_figure_plotly(figure, results_path / filename)
+    return figure
 
 
-def der_aggregated_active_power_time_series(results: data_models.RunResults, results_path: pathlib.Path):
+def der_aggregated_active_power_time_series(figure: go.Figure, results: data_models.RunResults) -> go.Figure:
     title = f"{constants.ValueLabels.ACTIVE_POWER} aggregated for all DERs"
-    filename = der_active_power_time_series.__name__
     x_label = constants.ValueLabels.TIME
     y_label = f"{constants.ValueLabels.ACTIVE_POWER} [{constants.ValueUnitLabels.WATT}]"
     line_name = constants.ValueLabels.DERS
 
-    figure = go.Figure()
     values = results.der_operation_results.der_active_power_vector.sum(axis="columns")
     figure.add_trace(go.Scatter(x=values.index, y=values.values, name=line_name))
     figure.update_layout(
@@ -88,17 +80,15 @@ def der_aggregated_active_power_time_series(results: data_models.RunResults, res
         yaxis_title=y_label,
         showlegend=False,
     )
-    plot_utils.write_figure_plotly(figure, results_path / filename)
+    return figure
 
 
-def der_aggregated_reactive_power_time_series(results: data_models.RunResults, results_path: pathlib.Path):
+def der_aggregated_reactive_power_time_series(figure: go.Figure, results: data_models.RunResults) -> go.Figure:
     title = f"{constants.ValueLabels.REACTIVE_POWER} aggregated for all DERs"
-    filename = der_reactive_power_time_series.__name__
     x_label = constants.ValueLabels.TIME
     y_label = f"{constants.ValueLabels.REACTIVE_POWER} [{constants.ValueUnitLabels.VOLT_AMPERE_REACTIVE}]"
     line_name = constants.ValueLabels.DERS
 
-    figure = go.Figure()
     values = results.der_operation_results.der_reactive_power_vector.sum(axis="columns")
     figure.add_trace(go.Scatter(x=values.index, y=values.values, name=line_name))
     figure.update_layout(
@@ -107,17 +97,15 @@ def der_aggregated_reactive_power_time_series(results: data_models.RunResults, r
         yaxis_title=y_label,
         showlegend=False,
     )
-    plot_utils.write_figure_plotly(figure, results_path / filename)
+    return figure
 
 
-def der_aggregated_apparent_power_time_series(results: data_models.RunResults, results_path: pathlib.Path):
+def der_aggregated_apparent_power_time_series(figure: go.Figure, results: data_models.RunResults) -> go.Figure:
     title = f"{constants.ValueLabels.APPARENT_POWER} aggregated for all DERs"
-    filename = der_apparent_power_time_series.__name__
     x_label = constants.ValueLabels.TIME
     y_label = f"{constants.ValueLabels.APPARENT_POWER} [{constants.ValueUnitLabels.VOLT_AMPERE}]"
     line_name = constants.ValueLabels.DERS
 
-    figure = go.Figure()
     # TODO: Add apparent power in result directly
     values = np.sqrt(
         results.der_operation_results.der_active_power_vector.sum(axis="columns") ** 2
@@ -130,35 +118,34 @@ def der_aggregated_apparent_power_time_series(results: data_models.RunResults, r
         yaxis_title=y_label,
         showlegend=False,
     )
-    plot_utils.write_figure_plotly(figure, results_path / filename)
+    return figure
+
 
-def node_voltage_per_unit_time_series(results: data_models.RunResults, results_path: pathlib.Path):
+def node_voltage_per_unit_time_series(figure: go.Figure, results: data_models.RunResults) -> go.Figure:
     title = f"{constants.ValueLabels.VOLTAGE} per Nodes"
-    filename = node_voltage_per_unit_time_series.__name__
     x_label = constants.ValueLabels.TIME
     y_label = f"{constants.ValueLabels.VOLTAGE} [{constants.ValueUnitLabels.VOLT_PER_UNIT}]"
     legend_title = constants.ValueLabels.NODES
 
-    figure = go.Figure()
     for node_type, node_name, phase in results.electric_grid_model_index.nodes:
-        values = results.electric_grid_operation_results.node_voltage_magnitude_vector_per_unit.loc[:, (slice(None), node_name, slice(None))].mean(axis="columns")
+        values = results.electric_grid_operation_results.node_voltage_magnitude_vector_per_unit.loc[
+            :, (slice(None), node_name, slice(None))
+        ].mean(axis="columns")
         figure.add_trace(go.Scatter(x=values.index, y=values.values, name=f"{node_name} ({node_type})"))
     figure.update_layout(
         title=title,
         xaxis_title=x_label,
         yaxis_title=y_label,
         legend=go.layout.Legend(title=legend_title, x=0.99, xanchor="auto", y=0.99, yanchor="auto"),
     )
-    plot_utils.write_figure_plotly(figure, results_path / filename)
+    return figure
 
 
-def node_aggregated_voltage_per_unit_time_series(results: data_models.RunResults, results_path: pathlib.Path):
+def node_aggregated_voltage_per_unit_time_series(figure: go.Figure, results: data_models.RunResults) -> go.Figure:
     title = f"{constants.ValueLabels.VOLTAGE} aggregated for all Nodes"
-    filename = node_voltage_per_unit_time_series.__name__
     x_label = constants.ValueLabels.TIME
     y_label = f"{constants.ValueLabels.VOLTAGE} [{constants.ValueUnitLabels.VOLT_PER_UNIT}]"
 
-    figure = go.Figure()
     for timestep in results.electric_grid_model_index.timesteps:
         values = results.electric_grid_operation_results.node_voltage_magnitude_vector_per_unit.loc[timestep, :]
         figure.add_trace(go.Box(name=timestep.isoformat(), y=values.T.values))
@@ -168,4 +155,4 @@ def node_aggregated_voltage_per_unit_time_series(results: data_models.RunResults
         yaxis_title=y_label,
         showlegend=False,
     )
-    plot_utils.write_figure_plotly(figure, results_path / filename)
+    return figure