feat: better results saving

pyproject.toml

@@ -30,6 +30,7 @@ numpy = "^1.20.1"
 jsonpickle = "^2.0.0"
 graphviz = "^0.16"
 networkx = "^2.5.1"
+matplotlib = "^3.4.1"
 
 [tool.poetry.dev-dependencies]
 pytest = "^6.2.2"

quaentropy/api/data_reader.py

@@ -5,7 +5,7 @@
 
 from pandas import DataFrame
 
-from quaentropy.api.data_writer import PlotDataType, BokehPlotGenerator
+from quaentropy.api.data_writer import PlotGenerator
 
 
 class ScriptViewer:
@@ -68,8 +68,7 @@ class PlotRecord:
     experiment_id: int
     id: int
     plot_data: Any = None
-    data_type: PlotDataType = PlotDataType.unknown
-    bokeh_generator: Optional[BokehPlotGenerator] = None
+    generator: Optional[PlotGenerator] = None
     label: Optional[str] = None
     story: Optional[str] = None
 

quaentropy/api/data_writer.py

@@ -1,11 +1,11 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from datetime import datetime
-from enum import Enum
-from typing import Any, Optional
+from typing import Any, Optional, Type
 
 from bokeh.models import Renderer
 from bokeh.plotting import Figure
+from matplotlib.figure import Figure as matplotlibFigure
 
 
 @dataclass
@@ -47,28 +47,22 @@ class Debug:
     extra: str
 
 
-class PlotDataType(Enum):
-    unknown = 0
-    np_2d = 1
-    py_2d = 2
-
-
-class BokehPlotGenerator(ABC):
+class PlotGenerator(ABC):
     def __init__(self) -> None:
         super().__init__()
 
     @abstractmethod
-    def plot_in_figure(
-        self, figure: Figure, data, data_type: PlotDataType, **kwargs
-    ) -> Renderer:
+    def plot_bokeh(self, figure: Figure, data, **kwargs) -> Renderer:
         pass
 
+    @abstractmethod
+    def plot_matplotlib(self, figure: matplotlibFigure, data, **kwargs):
+        pass
 
-@dataclass
-class Plot:
-    data: Any
-    data_type: PlotDataType
-    bokeh_generator: Optional[BokehPlotGenerator] = None
+
+@dataclass(frozen=True, eq=True)
+class PlotSpec:
+    generator: Optional[Type[PlotGenerator]] = None
     label: Optional[str] = None
     story: Optional[str] = ""
 
@@ -78,6 +72,7 @@ class NodeData:
     node_id: int
     start_time: datetime
     label: str
+    is_key_node: bool
 
 
 class DataWriter(ABC):
@@ -105,7 +100,7 @@ def save_debug(self, experiment_id: int, debug: Debug):
         pass
 
     @abstractmethod
-    def save_plot(self, experiment_id: int, plot: Plot):
+    def save_plot(self, experiment_id: int, plot: PlotSpec, data: Any):
         pass
 
     @abstractmethod

quaentropy/api/execution.py

@@ -6,7 +6,7 @@
     DataWriter,
     RawResultData,
     Metadata,
-    Plot,
+    PlotSpec,
 )
 from quaentropy.instruments.lab_topology import ExperimentResources
 
@@ -29,8 +29,8 @@ def add_result(self, result: RawResultData):
     def add_metadata(self, metadata: Metadata):
         self._data_writer.save_metadata(self._exp_id, metadata)
 
-    def add_plot(self, plot: Plot):
-        self._data_writer.save_plot(self._exp_id, plot)
+    def add_plot(self, plot: PlotSpec, data: Any):
+        self._data_writer.save_plot(self._exp_id, plot, data)
 
     def get_resource(self, name):
         return self._used_topology.get_resource(name)

quaentropy/api/graph.py

@@ -57,12 +57,15 @@ def get_parents(self) -> List[Node]:
         )
 
     def get_inputs(self) -> List[Output]:
-        return [var for var in self._input_vars.values()]
+        return list(self._input_vars.values())
+
+    def get_inputs_by_name(self) -> Dict[str, Output]:
+        return self._input_vars
 
     @abstractmethod
     async def execute_async(
         self,
-        parents_results: List[Dict[str, Any]],
+        input_values: Dict[str, Any],
         context: EntropyContext,
         node_execution_id: int,
         is_last,
@@ -73,7 +76,7 @@ async def execute_async(
     @abstractmethod
     def execute(
         self,
-        parents_results: List[Dict[Output, Any]],
+        input_values: Dict[str, Any],
         context: EntropyContext,
         node_execution_id: int,
         is_last,
@@ -95,15 +98,15 @@ def __init__(
         self,
         nodes: Set[Node],
         label: str = None,
-        plot_outputs: Optional[Set[str]] = None,
+        key_nodes: Optional[Set[Node]] = None,
     ) -> None:
         super().__init__()
         self.label: str = label
         self._nodes: Set[Node] = nodes
         self._end_nodes = self._calculate_last_nodes(nodes)
-        self._plot_outputs = plot_outputs
-        if self._plot_outputs is None:
-            self._plot_outputs = set()
+        self._key_nodes = key_nodes
+        if self._key_nodes is None:
+            self._key_nodes = set()
 
     def _calculate_last_nodes(self, nodes):
         end_nodes = nodes.copy()
@@ -117,12 +120,13 @@ def _calculate_last_nodes(self, nodes):
     def nodes(self) -> Set[Node]:
         return self._nodes
 
+    @property
     def end_nodes(self) -> List[Node]:
         return self._end_nodes
 
     @property
-    def plot_outputs(self) -> Set[str]:
-        return self._plot_outputs
+    def key_nodes(self) -> Set[Node]:
+        return self._key_nodes
 
     def export_dot_graph(self) -> Digraph:
         """

quaentropy/api/memory_reader_writer.py

@@ -1,6 +1,6 @@
 from datetime import datetime
 from time import time_ns
-from typing import List, Optional, Iterable
+from typing import List, Optional, Iterable, Any, Dict
 
 from pandas import DataFrame
 
@@ -13,7 +13,7 @@
     ScriptViewer,
     PlotRecord,
 )
-from quaentropy.api.data_writer import DataWriter, Plot, NodeData
+from quaentropy.api.data_writer import DataWriter, PlotSpec, NodeData
 from quaentropy.api.data_writer import (
     ExperimentInitialData,
     ExperimentEndData,
@@ -32,7 +32,7 @@ def __init__(self):
         self._results: List[RawResultData] = []
         self._metadata: List[Metadata] = []
         self._debug: Optional[Debug] = None
-        self._plot: List[Plot] = []
+        self._plot: Dict[PlotSpec, Any] = {}
         self._nodes: List[NodeData] = []
 
     def save_experiment_initial_data(self, initial_data: ExperimentInitialData) -> int:
@@ -51,8 +51,8 @@ def save_metadata(self, experiment_id: int, metadata: Metadata):
     def save_debug(self, experiment_id: int, debug: Debug):
         self._debug = debug
 
-    def save_plot(self, experiment_id: int, plot: Plot):
-        self._plot.append(plot)
+    def save_plot(self, experiment_id: int, plot: PlotSpec, data: Any):
+        self._plot[plot] = data
 
     def save_node(self, experiment_id: int, node_data: NodeData):
         self._nodes.append(node_data)
@@ -137,10 +137,9 @@ def get_plots(self, experiment_id: int) -> List[PlotRecord]:
         return [
             PlotRecord(
                 experiment_id,
-                self._plot.index(plot),
-                plot.data,
-                plot.data_type,
-                plot.bokeh_generator,
+                id(plot),
+                self._plot[plot],
+                plot.generator(),
                 plot.label,
                 plot.story,
             )