Add storage_graph to show the plugins stored or needed to be calcul…

…ated in the dependency tree (#1353)

* Add `storage_graph` to show the plugins stored or needed to be calculated in the dependency tree

* Add test of the new function

* Fix small bug and add more into `.gitignore`

* Try to install graphviz

* Be root

* Install graphviz

* Save run_id in the filename

* Add the color list to docstring

* update docstring

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix typo in docstring

* Minor change

---------

Co-authored-by: yuema137 <3124558229@qq.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

.github/workflows/pytest.yml

@@ -54,6 +54,9 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
 
+      - name: Install dependencies
+        run: sudo apt-get install -y graphviz
+
       - name: Install requirements
         run: |
           pip install -r extra_requirements/requirements-tests.txt

.github/workflows/test_install.yml

@@ -28,7 +28,9 @@ jobs:
           python-version: ${{ matrix.python-version }}
       - name: Checkout repo
         uses: actions/checkout@v4
-      - name: pre-install requirements
+      - name: Install dependencies
+        run: sudo apt-get install -y graphviz
+      - name: Pre-install requirements
         run: pip install -r requirements.txt
       - name: Install straxen
         run: python setup.py install

.gitignore

@@ -66,6 +66,12 @@ docs/_build
 # coverage
 .coverage
 
+# graphviz
+*.gv
+*.dot
+*.png
+*.svg
+
 # Documentation files
 docs/source/reference/datastructure*
 docs/source/reference/data_kinds*

requirements.txt

@@ -4,6 +4,7 @@
 bokeh
 commentjson
 gitpython
+graphviz
 immutabledict
 matplotlib
 multihist>=0.6.3

straxen/misc.py

@@ -3,16 +3,19 @@
 import inspect
 import warnings
 import datetime
+from immutabledict import immutabledict
 import pytz
 from itertools import chain
 from collections import defaultdict, OrderedDict, deque
 from importlib import import_module
 from platform import python_version
 import typing as ty
+from copy import deepcopy
 
 import numpy as np
 import pandas as pd
 import socket
+import graphviz
 import strax
 import straxen
 from git import Repo, InvalidGitRepositoryError
@@ -393,3 +396,87 @@ def sizeof(o):
         return s
 
     return sizeof(o)
+
+
+@strax.Context.add_method
+def storage_graph(
+    st,
+    run_id,
+    target,
+    graph=None,
+    not_stored=None,
+    dump_plot=True,
+    to_dir="./",
+    format="svg",
+):
+    """Plot the dependency graph indicating the storage of the plugins.
+
+    :param target: str of the target plugin to check
+    :param graph: graphviz.graphs.Digraph instance
+    :param not_stored: set of plugins which are not stored
+    :param dump_plot: bool, if True, save the plot to the to_dir
+    :param to_dir: str, directory to save the plot
+    :param format: str, format of the plot
+    :return: all plugins that will be calculated when running self.make(run_id, target)
+
+    The colors used in the graph represent the following storage states:
+    - grey: strax.SaveWhen.NEVER
+    - red: strax.SaveWhen.EXPLICIT
+    - orange: strax.SaveWhen.TARGET
+    - yellow: strax.SaveWhen.ALWAYS
+    - green: target is stored
+
+    """
+    save_when_colors = {
+        strax.SaveWhen.NEVER: "grey",
+        strax.SaveWhen.EXPLICIT: "red",
+        strax.SaveWhen.TARGET: "orange",
+        strax.SaveWhen.ALWAYS: "yellow",
+    }
+    if not_stored is None:
+        not_stored = set()
+    # the set of plugins which are not stored
+    if st.is_stored(run_id, target):
+        # if the plugin is stored, fill in green
+        fillcolor = "green"
+    else:
+        save_when = deepcopy(st._plugin_class_registry[target].save_when)
+        if isinstance(save_when, immutabledict):
+            save_when = save_when[target]
+        # if it is not stored, fill in the color according to save_when
+        fillcolor = save_when_colors[save_when]
+
+    if graph is None:
+        graph = graphviz.Digraph(name=f"{run_id}-{target}", strict=True)
+        graph.attr(bgcolor="transparent")
+    else:
+        if not isinstance(graph, graphviz.graphs.Digraph):
+            raise ValueError("graph should be a graphviz.Digraph instance!")
+    # add a node of target
+    graph.node(
+        target,
+        style="filled",
+        fillcolor=fillcolor,
+    )
+    if (not st.is_stored(run_id, target)) and (target not in not_stored):
+        not_stored.add(target)
+        depends_on = deepcopy(st._plugin_class_registry[target].depends_on)
+        depends_on = strax.to_str_tuple(depends_on)
+        for dep in depends_on:
+            # only add the node to the graph but not save the plot
+            not_stored.update(
+                st.storage_graph(
+                    run_id,
+                    dep,
+                    graph=graph,
+                    not_stored=not_stored,
+                    dump_plot=False,
+                    to_dir=to_dir,
+                )
+            )
+            graph.edge(target, dep)
+
+    # dump the plot if need
+    if dump_plot:
+        graph.render(format=format)
+    return not_stored

tests/test_basics.py

@@ -35,7 +35,7 @@ def test_run_selection(self):
         run_df = st.select_runs(available="raw_records")
         print(run_df)
         run_id = run_df.iloc[0]["name"]
-        assert run_id == test_run_id_1T
+        assert run_id == self.run_id
 
     def test_mini_analysis(self):
         @straxen.mini_analysis(requires=("raw_records",))
@@ -81,3 +81,7 @@ def test_raw_records_lineage(self):
         """The raw records lineage may NEVER change, if you ever do, doom ensures."""
         st = straxen.contexts.xenonnt_online()
         self.assertTrue(st.key_for("0", "raw_records").lineage_hash == "rfzvpzj4mf")
+
+    def test_storage_graph(self):
+        """Test the storage graph."""
+        self.st.storage_graph(self.run_id, "event_info")