Verified Commit e41ef748 authored by Yunus Sevinchan's avatar Yunus Sevinchan
Browse files

Add `max_num_nodes` option to DAG visualization 

Allows to not plot large DAGs, which often takes up long times
parent 2d7b07a2

dantro/plot/creators/base.py

+20 −4
Original line number Diff line number Diff line
@@ -908,6 +908,7 @@ class BasePlotCreator(AbstractPlotCreator): 
        output: dict = None,

        export: dict = None,

        generation: dict = None,

        max_num_nodes: int = 256,

        **plot_kwargs,

    ) -> Union["networkx.DiGraph", None]:

        """Generates a DAG representation according to certain criteria and
@@ -982,6 +983,9 @@ class BasePlotCreator(AbstractPlotCreator): 


            generation (dict, optional): Graph generation arguments passed to

                :py:meth:`~dantro.dag.TransformationDAG.generate_nx_graph`.

            max_num_nodes (int, optional): If given, do not continue with the

                visualization if the generated graph has more than this amount

                of nodes. Exporting is not affected by this.

            **plot_kwargs: Plotting-related arguments, passed on to

                :py:meth:`~dantro.dag.TransformationDAG.visualize`.
@@ -1120,16 +1124,28 @@ class BasePlotCreator(AbstractPlotCreator): 
            with exception_handling("exporting DAG representation"):

                export_graph(g, out_path=out_path, **export)



        # Plot it

        # Plot it (if there are not too many nodes)

        if plot_enabled:

            if max_num_nodes and g.number_of_nodes() > max_num_nodes:

                log.caution(

                    "Not visualizing this DAG because there are %d nodes in "

                    "the visualization, exceeding the configured limit of %d.",

                    g.number_of_nodes(),

                    max_num_nodes,

                )

                log.remark(

                    "Unset `max_num_nodes` or provide a larger limit to still "

                    "perform the visualization."

                )



            else:

                with exception_handling("plotting DAG representation"):

                    self._dag.visualize(g=g, out_path=out_path, **plot_kwargs)



                    if "error" in scenario:

                        log.caution(

                            "Created DAG visualization for scenario '%s'. "

                        "For debugging, inspecting the generated plot and the "

                        "traceback information may be helpful:\n  %s",

                            "For debugging, inspecting the generated plot and "

                            "the traceback information may be helpful:\n  %s",

                            scenario,

                            out_path,

                        )

tests/cfg/dag_plots.yml

+13 −0
Original line number Diff line number Diff line
@@ -391,6 +391,15 @@ config_based: 
          on_plot_error: true

          on_plot_success: true



  gen_nx_graph_max_num_nodes:

    plot_cfg:

      <<: *simple_example



      dag_visualization:

        when:

          always: true

        max_num_nodes: 3  # -> should lead to no plot being created



  gen_nx_graph_full_syntax:

    # NOTE does not use default values but aims to hit many branches

    plot_cfg:
@@ -917,6 +926,10 @@ config_based: 
        # exporting failed. If None, will use the creator's setting.

        raise_exc: ~



        # Maximum number of nodes for visualization; if set, will not perform

        # the plot if the number of nodes is exceeding this threshold.

        max_num_nodes: ~



        # Whether to *additionally* export the graph

        export:

          # Manipulate node or edge attributes (for export only)