lsst-pipe-base 29.2025.3900__py3-none-any.whl → 29.2025.4000__py3-none-any.whl

lsst/pipe/base/dot_tools.py

@@ -33,147 +33,27 @@ from __future__ import annotations
 
 __all__ = ["graph2dot", "pipeline2dot"]
 
-# -------------------------------
-#  Imports of standard modules --
-# -------------------------------
-import html
-import io
 from collections.abc import Iterable
 from typing import TYPE_CHECKING, Any
 
-# -----------------------------
-#  Imports for other modules --
-# -----------------------------
 from .pipeline import Pipeline
 
 if TYPE_CHECKING:
-    from lsst.daf.butler import DatasetRef
-    from lsst.pipe.base import QuantumGraph, QuantumNode, TaskDef
+    from .graph import QuantumGraph
+    from .pipeline import TaskDef
+    from .quantum_graph import PredictedQuantumGraph
 
-# ----------------------------------
-#  Local non-exported definitions --
-# ----------------------------------
 
-# Attributes applied to directed graph objects.
-_NODELABELPOINTSIZE = "18"
-_ATTRIBS = dict(
-    defaultGraph=dict(splines="ortho", nodesep="0.5", ranksep="0.75", pad="0.5"),
-    defaultNode=dict(shape="box", fontname="Monospace", fontsize="14", margin="0.2,0.1", penwidth="3"),
-    defaultEdge=dict(color="black", arrowsize="1.5", penwidth="1.5"),
-    task=dict(style="filled", color="black", fillcolor="#B1F2EF"),
-    quantum=dict(style="filled", color="black", fillcolor="#B1F2EF"),
-    dsType=dict(style="rounded,filled,bold", color="#00BABC", fillcolor="#F5F5F5"),
-    dataset=dict(style="rounded,filled,bold", color="#00BABC", fillcolor="#F5F5F5"),
-)
-
-
-def _renderDefault(type: str, attribs: dict[str, str], file: io.TextIOBase) -> None:
-    """Set default attributes for a given type."""
-    default_attribs = ", ".join([f'{key}="{val}"' for key, val in attribs.items()])
-    print(f"{type} [{default_attribs}];", file=file)
-
-
-def _renderNode(file: io.TextIOBase, nodeName: str, style: str, labels: list[str]) -> None:
-    """Render GV node"""
-    label = r"</TD></TR><TR><TD>".join(labels)
-    attrib_dict = dict(_ATTRIBS[style], label=label)
-    pre = '<<TABLE BORDER="0" CELLPADDING="5"><TR><TD>'
-    post = "</TD></TR></TABLE>>"
-    attrib = ", ".join(
-        [
-            f'{key}="{val}"' if key != "label" else f"{key}={pre}{val}{post}"
-            for key, val in attrib_dict.items()
-        ]
-    )
-    print(f'"{nodeName}" [{attrib}];', file=file)
-
-
-def _renderTaskNode(nodeName: str, taskDef: TaskDef, file: io.TextIOBase, idx: Any = None) -> None:
-    """Render GV node for a task"""
-    labels = [
-        f'<B><FONT POINT-SIZE="{_NODELABELPOINTSIZE}">' + html.escape(taskDef.label) + "</FONT></B>",
-        html.escape(taskDef.taskName),
-    ]
-    if idx is not None:
-        labels.append(f"<I>index:</I>&nbsp;{idx}")
-    if taskDef.connections:
-        # don't print collection of str directly to avoid visually noisy quotes
-        dimensions_str = ", ".join(sorted(taskDef.connections.dimensions))
-        labels.append(f"<I>dimensions:</I>&nbsp;{html.escape(dimensions_str)}")
-    _renderNode(file, nodeName, "task", labels)
-
-
-def _renderQuantumNode(
-    nodeName: str, taskDef: TaskDef, quantumNode: QuantumNode, file: io.TextIOBase
-) -> None:
-    """Render GV node for a quantum"""
-    labels = [f"{quantumNode.nodeId}", html.escape(taskDef.label)]
-    dataId = quantumNode.quantum.dataId
-    assert dataId is not None, "Quantum DataId cannot be None"
-    labels.extend(f"{key} = {dataId[key]}" for key in sorted(dataId.required.keys()))
-    _renderNode(file, nodeName, "quantum", labels)
-
-
-def _renderDSTypeNode(name: str, dimensions: list[str], file: io.TextIOBase) -> None:
-    """Render GV node for a dataset type"""
-    labels = [f'<B><FONT POINT-SIZE="{_NODELABELPOINTSIZE}">' + html.escape(name) + "</FONT></B>"]
-    if dimensions:
-        labels.append("<I>dimensions:</I>&nbsp;" + html.escape(", ".join(sorted(dimensions))))
-    _renderNode(file, name, "dsType", labels)
-
-
-def _renderDSNode(nodeName: str, dsRef: DatasetRef, file: io.TextIOBase) -> None:
-    """Render GV node for a dataset"""
-    labels = [html.escape(dsRef.datasetType.name), f"run: {dsRef.run!r}"]
-    labels.extend(f"{key} = {dsRef.dataId[key]}" for key in sorted(dsRef.dataId.required.keys()))
-    _renderNode(file, nodeName, "dataset", labels)
-
-
-def _renderEdge(fromName: str, toName: str, file: io.TextIOBase, **kwargs: Any) -> None:
-    """Render GV edge"""
-    if kwargs:
-        attrib = ", ".join([f'{key}="{val}"' for key, val in kwargs.items()])
-        print(f'"{fromName}" -> "{toName}" [{attrib}];', file=file)
-    else:
-        print(f'"{fromName}" -> "{toName}";', file=file)
-
-
-def _datasetRefId(dsRef: DatasetRef) -> str:
-    """Make an identifying string for given ref"""
-    dsId = [dsRef.datasetType.name]
-    dsId.extend(f"{key} = {dsRef.dataId[key]}" for key in sorted(dsRef.dataId.required.keys()))
-    return ":".join(dsId)
-
-
-def _makeDSNode(dsRef: DatasetRef, allDatasetRefs: dict[str, str], file: io.TextIOBase) -> str:
-    """Make new node for dataset if  it does not exist.
-
-    Returns node name.
-    """
-    dsRefId = _datasetRefId(dsRef)
-    nodeName = allDatasetRefs.get(dsRefId)
-    if nodeName is None:
-        idx = len(allDatasetRefs)
-        nodeName = f"dsref_{idx}"
-        allDatasetRefs[dsRefId] = nodeName
-        _renderDSNode(nodeName, dsRef, file)
-    return nodeName
-
-
-# ------------------------
-#  Exported definitions --
-# ------------------------
-
-
-def graph2dot(qgraph: QuantumGraph, file: Any) -> None:
+def graph2dot(qgraph: QuantumGraph | PredictedQuantumGraph, file: Any) -> None:
     """Convert QuantumGraph into GraphViz digraph.
 
     This method is mostly for documentation/presentation purposes.
 
     Parameters
     ----------
-    qgraph : `lsst.pipe.base.QuantumGraph`
-        QuantumGraph instance.
+    qgraph : `lsst.pipe.base.QuantumGraph` or \
+            `lsst.pipe.base.quantum_graph.PredictedQuantumGraph`
+        Quantum graph object.
     file : `str` or file object
         File where GraphViz graph (DOT language) is written, can be a file name
         or file object.
@@ -185,38 +65,20 @@ def graph2dot(qgraph: QuantumGraph, file: Any) -> None:
     ImportError
         Raised if the task class cannot be imported.
     """
+    from .quantum_graph import PredictedQuantumGraph, visualization
+
+    if not isinstance(qgraph, PredictedQuantumGraph):
+        qgraph = PredictedQuantumGraph.from_old_quantum_graph(qgraph)
+
     # open a file if needed
     close = False
     if not hasattr(file, "write"):
         file = open(file, "w")
         close = True
 
-    print("digraph QuantumGraph {", file=file)
-    _renderDefault("graph", _ATTRIBS["defaultGraph"], file)
-    _renderDefault("node", _ATTRIBS["defaultNode"], file)
-    _renderDefault("edge", _ATTRIBS["defaultEdge"], file)
-
-    allDatasetRefs: dict[str, str] = {}
-    for taskId, taskDef in enumerate(qgraph.taskGraph):
-        quanta = qgraph.getNodesForTask(taskDef)
-        for qId, quantumNode in enumerate(quanta):
-            # node for a task
-            taskNodeName = f"task_{taskId}_{qId}"
-            _renderQuantumNode(taskNodeName, taskDef, quantumNode, file)
-
-            # quantum inputs
-            for dsRefs in quantumNode.quantum.inputs.values():
-                for dsRef in dsRefs:
-                    nodeName = _makeDSNode(dsRef, allDatasetRefs, file)
-                    _renderEdge(nodeName, taskNodeName, file)
-
-            # quantum outputs
-            for dsRefs in quantumNode.quantum.outputs.values():
-                for dsRef in dsRefs:
-                    nodeName = _makeDSNode(dsRef, allDatasetRefs, file)
-                    _renderEdge(taskNodeName, nodeName, file)
+    v = visualization.QuantumGraphDotVisualizer()
+    v.write_bipartite(qgraph, file)
 
-    print("}", file=file)
     if close:
         file.close()
 

lsst/pipe/base/exec_fixup_data_id.py

@@ -27,16 +27,17 @@
 
 __all__ = ["ExecutionGraphFixup"]
 
-import contextlib
 import itertools
+import uuid
 from collections import defaultdict
-from collections.abc import Sequence
-from typing import Any
+from collections.abc import Mapping, Sequence
 
 import networkx as nx
 
+from lsst.daf.butler import DataCoordinate, DataIdValue
+
 from .execution_graph_fixup import ExecutionGraphFixup
-from .graph import QuantumGraph, QuantumNode
+from .graph import QuantumGraph
 
 
 class ExecFixupDataId(ExecutionGraphFixup):
@@ -88,44 +89,16 @@ class ExecFixupDataId(ExecutionGraphFixup):
         else:
             self.dimensions = tuple(self.dimensions)
 
-    def _key(self, qnode: QuantumNode) -> tuple[Any, ...]:
-        """Produce comparison key for quantum data.
-
-        Parameters
-        ----------
-        qnode : `QuantumNode`
-            An individual node in a `~lsst.pipe.base.QuantumGraph`
-
-        Returns
-        -------
-        key : `tuple`
-        """
-        dataId = qnode.quantum.dataId
-        assert dataId is not None, "Quantum DataId cannot be None"
-        key = tuple(dataId[dim] for dim in self.dimensions)
-        return key
-
     def fixupQuanta(self, graph: QuantumGraph) -> QuantumGraph:
-        taskDef = graph.findTaskDefByLabel(self.taskLabel)
-        if taskDef is None:
-            raise ValueError(f"Cannot find task with label {self.taskLabel}")
-        quanta = list(graph.getNodesForTask(taskDef))
-        keyQuanta = defaultdict(list)
-        for q in quanta:
-            key = self._key(q)
-            keyQuanta[key].append(q)
-        keys = sorted(keyQuanta.keys(), reverse=self.reverse)
-        networkGraph = graph.graph
-
-        for prev_key, key in itertools.pairwise(keys):
-            for prev_node in keyQuanta[prev_key]:
-                for node in keyQuanta[key]:
-                    # remove any existing edges between the two nodes, but
-                    # don't fail if there are not any. Both directions need
-                    # tried because in a directed graph, order maters
-                    for edge in ((node, prev_node), (prev_node, node)):
-                        with contextlib.suppress(nx.NetworkXException):
-                            networkGraph.remove_edge(*edge)
-
-                    networkGraph.add_edge(prev_node, node)
-        return graph
+        raise NotImplementedError()
+
+    def fixup_graph(
+        self, xgraph: nx.DiGraph, quanta_by_task: Mapping[str, Mapping[DataCoordinate, uuid.UUID]]
+    ) -> None:
+        quanta_by_sort_key: defaultdict[tuple[DataIdValue, ...], list[uuid.UUID]] = defaultdict(list)
+        for data_id, quantum_id in quanta_by_task[self.taskLabel].items():
+            key = tuple(data_id[dim] for dim in self.dimensions)
+            quanta_by_sort_key[key].append(quantum_id)
+        sorted_keys = sorted(quanta_by_sort_key.keys(), reverse=self.reverse)
+        for prev_key, key in itertools.pairwise(sorted_keys):
+            xgraph.add_edges_from(itertools.product(quanta_by_sort_key[prev_key], quanta_by_sort_key[key]))

lsst/pipe/base/execution_graph_fixup.py

@@ -27,7 +27,13 @@
 
 __all__ = ["ExecutionGraphFixup"]
 
-from abc import ABC, abstractmethod
+import uuid
+from abc import ABC
+from collections.abc import Mapping
+
+import networkx
+
+from lsst.daf.butler import DataCoordinate
 
 from .graph import QuantumGraph
 
@@ -35,27 +41,26 @@ from .graph import QuantumGraph
 class ExecutionGraphFixup(ABC):
     """Interface for classes which update quantum graphs before execution.
 
-    Primary goal of this class is to modify quanta dependencies which may not
-    be possible to reflect in a quantum graph using standard tools. One known
-    use case for that is to guarantee particular execution order of visits in
-    CI jobs for cases when outcome depends on the processing order of visits
-    (e.g. AP association pipeline).
-
-    Instances of this class receive pre-ordered sequence of quanta
-    (`.QuantumGraph` instances) and they are allowed to modify quanta data in
-    place, for example update ``dependencies`` field to add additional
-    dependencies. Returned list of quanta will be re-ordered once again by the
-    graph executor to reflect new dependencies.
+    Notes
+    -----
+    The primary goal of this class is to modify quanta dependencies which may
+    not be possible to reflect in a quantum graph using standard tools. One
+    known use case for that is to guarantee particular execution order of
+    visits in CI jobs for cases when outcome depends on the processing order of
+    visits (e.g. AP association pipeline).
+
+    Instances of this class receive a preliminary graph and are allowed to
+    add edges, as long as those edges do not result in a cycle.  Edges and
+    nodes may not be removed.
+
+    New subclasses should implement only `fixup_graph`, which will always be
+    called first.  `fixupQuanta` is only called if `fixup_graph` raises
+    `NotImplementedError`.
     """
 
-    @abstractmethod
     def fixupQuanta(self, graph: QuantumGraph) -> QuantumGraph:
         """Update quanta in a graph.
 
-        Potentially anything in the graph could be changed if it does not
-        break executor assumptions. If modifications result in a dependency
-        cycle the executor will raise an exception.
-
         Parameters
         ----------
         graph : `.QuantumGraph`
@@ -65,5 +70,31 @@ class ExecutionGraphFixup(ABC):
         -------
         graph : `.QuantumGraph`
             Modified graph.
+
+        Notes
+        -----
+        This hook is provided for backwards compatibility only.
+        """
+        raise NotImplementedError()
+
+    def fixup_graph(
+        self, xgraph: networkx.DiGraph, quanta_by_task: Mapping[str, Mapping[DataCoordinate, uuid.UUID]]
+    ) -> None:
+        """Update a networkx graph of quanta in place by adding edges to
+        further constrain the ordering.
+
+        Parameters
+        ----------
+        xgraph : `networkx.DiGraph`
+            A directed acyclic graph of quanta to modify in place.  Node keys
+            are quantum UUIDs, and attributes include ``task_label`` (`str`)
+            and ``data_id`` (a full `lsst.daf.butler.DataCoordinate`, without
+            dimension records attached).  Edges may be added, but not removed.
+            Nodes may not be modified.
+        quanta_by_task : `~collections.abc.Mapping` [ `str`,\
+                `~collections.abc.Mapping` [ `lsst.daf.butler.DataCoordinate`,\
+                `uuid.UUID` ] ]
+            All quanta in the graph, grouped first by task label and then by
+            data ID.
         """
-        raise NotImplementedError
+        raise NotImplementedError()

lsst/pipe/base/graph/graph.py

@@ -806,11 +806,18 @@ class QuantumGraph:
         uri : convertible to `~lsst.resources.ResourcePath`
             URI to where the graph should be saved.
         """
-        buffer = self._buildSaveObject()
         path = ResourcePath(uri)
-        if path.getExtension() not in (".qgraph"):
-            raise TypeError(f"Can currently only save a graph in qgraph format not {uri}")
-        path.write(buffer)  # type: ignore  # Ignore because bytearray is safe to use in place of bytes
+        match path.getExtension():
+            case ".qgraph":
+                buffer = self._buildSaveObject()
+                path.write(buffer)  # type: ignore  # Ignore because bytearray is safe to use in place of bytes
+            case ".qg":
+                from ..quantum_graph import PredictedQuantumGraphComponents
+
+                pqg = PredictedQuantumGraphComponents.from_old_quantum_graph(self)
+                pqg.write(path)
+            case ext:
+                raise TypeError(f"Can currently only save a graph in .qgraph or .qg format, not {ext!r}.")
 
     @property
     def metadata(self) -> MappingProxyType[str, Any]:
@@ -962,11 +969,23 @@ class QuantumGraph:
             the graph.
         """
         uri = ResourcePath(uri)
-        if uri.getExtension() in {".qgraph"}:
-            with LoadHelper(uri, minimumVersion, fullRead=(nodes is None)) as loader:
-                qgraph = loader.load(universe, nodes, graphID)
-        else:
-            raise ValueError(f"Only know how to handle files saved as `.qgraph`, not {uri}")
+        match uri.getExtension():
+            case ".qgraph":
+                with LoadHelper(uri, minimumVersion, fullRead=(nodes is None)) as loader:
+                    qgraph = loader.load(universe, nodes, graphID)
+            case ".qg":
+                from ..quantum_graph import PredictedQuantumGraphReader
+
+                with PredictedQuantumGraphReader.open(uri, page_size=100000) as qgr:
+                    quantum_ids = (
+                        [uuid.UUID(q) if not isinstance(q, uuid.UUID) else q for q in nodes]
+                        if nodes is not None
+                        else None
+                    )
+                    qgr.read_execution_quanta(quantum_ids)
+                    qgraph = qgr.finish().to_old_quantum_graph()
+            case _:
+                raise ValueError(f"Only know how to handle files saved as `.qgraph`, not {uri}")
         if not isinstance(qgraph, QuantumGraph):
             raise TypeError(f"QuantumGraph file {uri} contains unexpected object type: {type(qgraph)}")
         return qgraph

lsst/pipe/base/graph_walker.py

@@ -0,0 +1,119 @@
+# This file is part of pipe_base.
+#
+# Developed for the LSST Data Management System.
+# This product includes software developed by the LSST Project
+# (http://www.lsst.org).
+# See the COPYRIGHT file at the top-level directory of this distribution
+# for details of code ownership.
+#
+# This software is dual licensed under the GNU General Public License and also
+# under a 3-clause BSD license. Recipients may choose which of these licenses
+# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
+# respectively.  If you choose the GPL option then the following text applies
+# (but note that there is still no warranty even if you opt for BSD instead):
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+from __future__ import annotations
+
+__all__ = ("GraphWalker",)
+
+from typing import Generic, Self, TypeVar
+
+import networkx
+
+_T = TypeVar("_T")
+
+
+class GraphWalker(Generic[_T]):
+    """A helper for traversing directed acyclic graphs.
+
+    Parameters
+    ----------
+    xgraph : `networkx.DiGraph` or `networkx.MultiDiGraph`
+        Networkx graph to process.  Will be consumed during iteration, so this
+        should often be a copy.
+
+    Notes
+    -----
+    Each iteration yields a `frozenset` of nodes, which may be empty if there
+    are no nodes ready for processing.  A node is only considered ready if all
+    of its predecessor nodes have been marked as complete with `finish`.
+    Iteration only completes when all nodes have been finished or failed.
+
+    `GraphWalker` is not thread-safe; calling one `GraphWalker` method while
+    another is in progress is undefined behavior.  It is designed to be used
+    in the management thread or process in a parallel traversal.
+    """
+
+    def __init__(self, xgraph: networkx.DiGraph | networkx.MultiDiGraph):
+        self._xgraph = xgraph
+        self._ready: set[_T] = set(next(iter(networkx.dag.topological_generations(self._xgraph)), []))
+        self._active: set[_T] = set()
+        self._incomplete: set[_T] = set(self._xgraph)
+
+    def __iter__(self) -> Self:
+        return self
+
+    def __next__(self) -> frozenset[_T]:
+        if not self._incomplete:
+            raise StopIteration()
+        new_active = frozenset(self._ready)
+        self._active.update(new_active)
+        self._ready.clear()
+        return new_active
+
+    def finish(self, key: _T) -> None:
+        """Mark a node as successfully processed, unblocking (at least in part)
+        iteration over successor nodes.
+
+        Parameters
+        ----------
+        key : unspecified
+            NetworkX key of the node to mark finished.
+        """
+        self._active.remove(key)
+        self._incomplete.remove(key)
+        successors = list(self._xgraph.successors(key))
+        for successor in successors:
+            assert successor not in self._active, (
+                "A node downstream of an active one should not have been yielded yet."
+            )
+            if all(
+                predecessor not in self._incomplete for predecessor in self._xgraph.predecessors(successor)
+            ):
+                self._ready.add(successor)
+
+    def fail(self, key: _T) -> list[_T]:
+        """Mark a node as unsuccessfully processed, permanently blocking all
+        recursive descendants.
+
+        Parameters
+        ----------
+        key : unspecified
+            NetworkX key of the node to mark as a failure.
+
+        Returns
+        -------
+        blocked : `list`
+            NetworkX keys of nodes that were recursive descendants of the
+            failed node, and will hence never be yielded by the iterator.
+        """
+        self._active.remove(key)
+        self._incomplete.remove(key)
+        descendants = list(networkx.dag.descendants(self._xgraph, key))
+        self._xgraph.remove_node(key)
+        self._xgraph.remove_nodes_from(descendants)
+        self._incomplete.difference_update(descendants)
+        return descendants

lsst/pipe/base/log_capture.py

@@ -41,6 +41,7 @@ from lsst.daf.butler import Butler, FileDataset, LimitedButler, Quantum
 from lsst.daf.butler.logging import ButlerLogRecordHandler, ButlerLogRecords, ButlerMDC, JsonLogFormatter
 
 from ._status import InvalidQuantumError
+from .automatic_connection_constants import METADATA_OUTPUT_TEMPLATE
 from .pipeline_graph import TaskNode
 
 _LOG = logging.getLogger(__name__)
@@ -116,8 +117,10 @@ class LogCapture:
         mdc = {"LABEL": task_node.label, "RUN": ""}
         if quantum.dataId:
             mdc["LABEL"] += f":{quantum.dataId}"
-        if self.full_butler is not None:
-            mdc["RUN"] = self.full_butler.run or ""
+
+        metadata_ref = quantum.outputs[METADATA_OUTPUT_TEMPLATE.format(label=task_node.label)][0]
+        mdc["RUN"] = metadata_ref.run
+
         ctx = _LogCaptureFlag()
         log_dataset_name = (
             task_node.log_output.dataset_type_name if task_node.log_output is not None else None

lsst/pipe/base/mermaid_tools.py

@@ -39,39 +39,12 @@ from typing import TYPE_CHECKING, Any, Literal
 from .pipeline import Pipeline
 
 if TYPE_CHECKING:
-    from lsst.daf.butler import DatasetRef
-    from lsst.pipe.base import QuantumGraph, TaskDef
+    from .graph import QuantumGraph
+    from .pipeline import TaskDef
+    from .quantum_graph import PredictedQuantumGraph
 
 
-def _datasetRefId(dsRef: DatasetRef) -> str:
-    """Make a unique identifier string for a dataset ref based on its name and
-    dataId.
-    """
-    dsIdParts = [dsRef.datasetType.name]
-    dsIdParts.extend(f"{key}_{dsRef.dataId[key]}" for key in sorted(dsRef.dataId.required.keys()))
-    return "_".join(dsIdParts)
-
-
-def _makeDatasetNode(dsRef: DatasetRef, allDatasetRefs: dict[str, str], file: Any) -> str:
-    """Create a Mermaid node for a dataset if it doesn't exist, and return its
-    node ID.
-    """
-    dsId = _datasetRefId(dsRef)
-    nodeName = allDatasetRefs.get(dsId)
-    if nodeName is None:
-        nodeName = f"DATASET_{len(allDatasetRefs)}"
-        allDatasetRefs[dsId] = nodeName
-        # Simple label: datasetType name and run.
-        label_lines = [f"**{dsRef.datasetType.name}**", f"run: {dsRef.run}"]
-        # Add dataId info.
-        for k in sorted(dsRef.dataId.required.keys()):
-            label_lines.append(f"{k}={dsRef.dataId[k]}")
-        label = "<br>".join(label_lines)
-        print(f'{nodeName}["{label}"]', file=file)
-    return nodeName
-
-
-def graph2mermaid(qgraph: QuantumGraph, file: Any) -> None:
+def graph2mermaid(qgraph: QuantumGraph | PredictedQuantumGraph, file: Any) -> None:
     """Convert QuantumGraph into a Mermaid flowchart (top-down).
 
     This method is mostly for documentation/presentation purposes.
@@ -91,45 +64,19 @@ def graph2mermaid(qgraph: QuantumGraph, file: Any) -> None:
     ImportError
         Raised if the task class cannot be imported.
     """
+    from .quantum_graph import PredictedQuantumGraph, visualization
+
+    if not isinstance(qgraph, PredictedQuantumGraph):
+        qgraph = PredictedQuantumGraph.from_old_quantum_graph(qgraph)
+
     # Open a file if needed.
     close = False
     if not hasattr(file, "write"):
         file = open(file, "w")
         close = True
 
-    # Start Mermaid code block with flowchart.
-    print("flowchart TD", file=file)
-
-    # To avoid duplicating dataset nodes, we track them.
-    allDatasetRefs: dict[str, str] = {}
-
-    # Process each task/quantum.
-    for taskId, taskDef in enumerate(qgraph.taskGraph):
-        quanta = qgraph.getNodesForTask(taskDef)
-        for qId, quantumNode in enumerate(quanta):
-            # Create quantum node.
-            taskNodeName = f"TASK_{taskId}_{qId}"
-            taskLabelLines = [f"**{taskDef.label}**", f"Node ID: {quantumNode.nodeId}"]
-            dataId = quantumNode.quantum.dataId
-            if dataId is not None:
-                for k in sorted(dataId.required.keys()):
-                    taskLabelLines.append(f"{k}={dataId[k]}")
-            else:
-                raise ValueError("Quantum DataId cannot be None")
-            taskLabel = "<br>".join(taskLabelLines)
-            print(f'{taskNodeName}["{taskLabel}"]', file=file)
-
-            # Quantum inputs: datasets --> tasks
-            for dsRefs in quantumNode.quantum.inputs.values():
-                for dsRef in dsRefs:
-                    dsNode = _makeDatasetNode(dsRef, allDatasetRefs, file)
-                    print(f"{dsNode} --> {taskNodeName}", file=file)
-
-            # Quantum outputs: tasks --> datasets
-            for dsRefs in quantumNode.quantum.outputs.values():
-                for dsRef in dsRefs:
-                    dsNode = _makeDatasetNode(dsRef, allDatasetRefs, file)
-                    print(f"{taskNodeName} --> {dsNode}", file=file)
+    v = visualization.QuantumGraphMermaidVisualizer()
+    v.write_bipartite(qgraph, file)
 
     if close:
         file.close()